longcat: the offload setting never did anything (v0.2.67)

LongCatVideoPipeline's docstring says "This model inherits from
[`DiffusionPipeline`]". The class declaration inherits nothing, so it has no
enable_model_cpu_offload and no enable_sequential_cpu_offload. The guarded
getattr found neither, logged a warning and fell through to pipe.to("cuda") —
so every load put all three components on the card and the offload mode in the
model config was decorative. That is the 45 GB ask against a 24 GB card.

An INT8 DiT does not fix it on its own: 13.6 GB of DiT plus ~11.2 GB for the
UMT5-XXL encoder (5.6B, stored fp32, loaded bf16) plus the VAE is ~25 GB before
activations.

But the encoder runs once, to encode the prompt, and is dead weight for every
denoising step after. So offload now keeps it on the host and brings it over
only for encode_prompt, set on the instance because all four generate_* methods
call self.encode_prompt. The DiT and VAE stay resident, which leaves the loop —
the part that has to be fast — on the card.

It is moved back in a finally: stranding 11 GB after a failed encode would make
the retry the thing that runs out of memory. 'sequential' is served as model
offload and says so, rather than implying a deeper mode that upstream has no
machinery for.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0185zEBiy3KWB37xNGSrxL7w

codai/__init__.py

@@ -16,7 +16,7 @@
 
 # Canonical product version for CoderAI — single source of truth. Both the API
 # metadata and the admin web UI read from here.
-__version__ = "0.2.66"
+__version__ = "0.2.67"
 __tagline__ = "Complete orchestration, distribution and escalation of remotizable advanced inference."
 
 # Configure the CUDA caching allocator BEFORE torch is imported anywhere.

tests/test_longcat_placement.py

+"""Where each component of the LongCat pipeline actually sits on the card.
+
+Upstream's LongCatVideoPipeline is a plain object whose docstring claims it inherits
+from DiffusionPipeline. It does not, so it has no enable_model_cpu_offload and no
+enable_sequential_cpu_offload: asking for offload logged a warning and then loaded
+everything onto the card anyway, and the offload setting in the model config did
+nothing. On a 24 GB card that is ~25 GB of components before activations even with an
+INT8 DiT, because the UMT5-XXL text encoder is ~11 GB on its own.
+
+The encoder is used once, to encode the prompt. These tests hold it off the card for
+everything else.
+"""
+import importlib.util
+import pathlib
+import sys
+
+import pytest
+
+ROOT = pathlib.Path(__file__).resolve().parents[1]
+SERVICE = ROOT / "tools" / "longcat_service.py"
+
+
+class FakeModule:
+    """Records every device it is moved to, in order."""
+
+    def __init__(self, name):
+        self.name = name
+        self.device = "cpu"
+        self.history = ["cpu"]
+
+    def to(self, device, non_blocking=False):
+        self.device = str(device)
+        self.history.append(str(device))
+        return self
+
+
+class FakePipeline:
+    """Upstream's shape: a plain object with its own to(), no offload hooks at all."""
+
+    def __init__(self):
+        self.dit = FakeModule("dit")
+        self.vae = FakeModule("vae")
+        self.text_encoder = FakeModule("text_encoder")
+        self.device = "cpu"
+        self.encode_calls = []
+        # What was on the card at the moment each component was moved there.
+        self.coresident = []
+
+    def to(self, device):
+        self.device = str(device)
+        for attr in ("dit", "text_encoder", "vae"):
+            mod = getattr(self, attr)
+            if mod is not None:
+                mod.to(device, non_blocking=True)
+        self._snapshot()
+        return self
+
+    def _snapshot(self):
+        self.coresident.append({
+            name: getattr(self, name).device
+            for name in ("dit", "vae", "text_encoder")
+            if getattr(self, name) is not None
+        })
+
+    def encode_prompt(self, prompt, device=None, **kwargs):
+        self.encode_calls.append(prompt)
+        self._snapshot()
+        return ("embeds", "mask")
+
+
+@pytest.fixture
+def service(monkeypatch):
+    torch = pytest.importorskip("torch")
+    # Never touch a real card from a test: production holds this one.
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None)
+    spec = importlib.util.spec_from_file_location("lc_service_placement", SERVICE)
+    mod = importlib.util.module_from_spec(spec)
+    monkeypatch.setitem(sys.modules, spec.name, mod)
+    spec.loader.exec_module(mod)
+    return mod
+
+
+def test_without_offload_everything_is_resident(service):
+    """The existing behaviour, for a card that can hold it all."""
+    pipe = FakePipeline()
+
+    service._place_pipeline(pipe, None)
+
+    assert pipe.dit.device == "cuda"
+    assert pipe.vae.device == "cuda"
+    assert pipe.text_encoder.device == "cuda"
+
+
+@pytest.mark.parametrize("mode", ["model", "sequential"])
+def test_offload_leaves_the_text_encoder_on_the_host(service, mode):
+    """~11 GB that is dead weight for every denoising step."""
+    pipe = FakePipeline()
+
+    service._place_pipeline(pipe, mode)
+
+    assert pipe.dit.device == "cuda", "the DiT must stay resident — it runs every step"
+    assert pipe.vae.device == "cuda"
+    assert pipe.text_encoder.device == "cpu"
+
+
+def test_the_encoder_is_never_on_the_card_beside_the_dit_during_placement(service):
+    """Moving it over and straight back would still need the peak we are avoiding."""
+    pipe = FakePipeline()
+
+    service._place_pipeline(pipe, "model")
+
+    for snapshot in pipe.coresident:
+        both = snapshot.get("dit") == "cuda" and snapshot.get("text_encoder") == "cuda"
+        assert not both, f"DiT and text encoder were both on the card: {snapshot}"
+    assert "cuda" not in pipe.text_encoder.history
+
+
+def test_the_encoder_visits_the_card_to_encode_and_leaves_again(service):
+    pipe = FakePipeline()
+    service._place_pipeline(pipe, "model")
+
+    out = pipe.encode_prompt("a fighter entering the ring", device="cuda")
+
+    assert out == ("embeds", "mask")
+    assert pipe.encode_calls == ["a fighter entering the ring"]
+    during = pipe.coresident[-1]
+    assert during["text_encoder"] == "cuda", "it has to be on the card to encode"
+    assert pipe.text_encoder.device == "cpu", "and has to go back afterwards"
+
+
+def test_a_failed_encode_still_puts_the_encoder_back(service):
+    """Stranding 11 GB on a failure makes the retry the thing that runs out of memory."""
+
+    class Failing(FakePipeline):
+        def encode_prompt(self, prompt, device=None, **kwargs):
+            self._snapshot()
+            raise RuntimeError("tokenizer blew up")
+
+    pipe = Failing()
+    service._place_pipeline(pipe, "model")
+
+    with pytest.raises(RuntimeError, match="tokenizer blew up"):
+        pipe.encode_prompt("x")
+
+    assert pipe.coresident[-1]["text_encoder"] == "cuda", "it did reach the card"
+    assert pipe.text_encoder.device == "cpu", "and was put back despite the failure"
+
+
+def test_every_generate_path_goes_through_the_wrapper(service):
+    """All four generate_* methods call self.encode_prompt, so the wrapper must be set
+    on the instance, where it shadows the class method."""
+    pipe = FakePipeline()
+    service._place_pipeline(pipe, "model")
+
+    assert "encode_prompt" in pipe.__dict__
+    assert pipe.encode_prompt is not FakePipeline.encode_prompt
+
+
+def test_no_cuda_is_reported_not_silently_tolerated(service, monkeypatch):
+    import torch
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+    pipe = FakePipeline()
+
+    service._place_pipeline(pipe, "model")
+
+    assert pipe.dit.device == "cpu"

tools/longcat_service.py

@@ -181,29 +181,79 @@ def load_pipeline(checkpoint_dir: str, source_dir: str, dtype: str = "bfloat16",
                 raise RuntimeError(
                     "use_distill needs the DiT's enable_loras() — this build of "
                     "longcat_video does not expose it")
-        # Offload is opt-in: the 13.6B DiT wants ~27 GB at bf16 and the reported peak for
-        # a full profile is ~41.6 GB, so a smaller card needs it. Not every build of the
-        # upstream pipeline exposes the diffusers hooks, hence the guarded calls.
-        placed = False
-        if offload in ("model", "sequential"):
-            fn = getattr(pipe, f"enable_{offload}_cpu_offload", None)
-            if callable(fn):
-                fn()
-                placed = True
-                log(f"{offload} CPU offload enabled")
-            else:
-                log(f"WARNING: this pipeline has no enable_{offload}_cpu_offload(); "
-                    f"loading fully on device instead")
-        if not placed:
-            if torch.cuda.is_available():
-                pipe.to("cuda")
-            else:
-                log("WARNING: no CUDA device visible — this will be extremely slow")
+        _place_pipeline(pipe, offload)
         _state.update(pipe=pipe, model=root, source=src, dtype=dtype)
         log("ready")
         return pipe
 
 
+def _place_pipeline(pipe, offload):
+    """Put the pipeline on the card, honouring the offload mode.
+
+    Upstream's LongCatVideoPipeline is a PLAIN OBJECT. Its docstring says "This model
+    inherits from [`DiffusionPipeline`]", but the class declaration inherits nothing, so
+    enable_model_cpu_offload() and enable_sequential_cpu_offload() do not exist on it.
+    Asking for either used to log a warning and then put everything on the card anyway —
+    which is how a 24 GB card came to be asked for ~45 GB, with the offload setting in
+    the model config doing nothing at all.
+
+    Full residency does not fit 24 GB even with an INT8 DiT: 13.6 GB of DiT, ~11.2 GB for
+    the UMT5-XXL encoder (5.6B, stored fp32, loaded bf16) and the VAE come to ~25 GB
+    before a single activation. But the encoder runs ONCE, to turn the prompt into
+    embeddings, and is dead weight for every denoising step after that. So offload keeps
+    it on the CPU and brings it over only for encode_prompt, which leaves the DiT — the
+    part that has to be fast — resident for the whole loop.
+    """
+    import gc
+    import torch
+
+    if not torch.cuda.is_available():
+        log("WARNING: no CUDA device visible — this will be extremely slow")
+        return pipe
+
+    if offload not in ("model", "sequential"):
+        pipe.to("cuda")
+        log("the whole pipeline is resident on the card")
+        return pipe
+    if offload == "sequential":
+        # Upstream has no per-submodule streaming to drive, and writing one would buy
+        # nothing here: the encoder is the only component that does not fit beside the
+        # DiT. Say what is actually being done rather than imply a deeper mode.
+        log("note: 'sequential' is served as model offload — the text encoder is the "
+            "only component that has to leave the card")
+
+    # Move everything EXCEPT the encoder, through upstream's own to() so the DiT's LoRA
+    # networks travel with it. Detaching the encoder first means it is never on the card
+    # and in host memory at the same time, which is the peak that has to be avoided.
+    encoder = pipe.text_encoder
+    pipe.text_encoder = None
+    try:
+        pipe.to("cuda")
+    finally:
+        pipe.text_encoder = encoder
+    if encoder is not None:
+        pipe.text_encoder = encoder.to("cpu")
+
+    inner = pipe.encode_prompt      # the bound method, before it is shadowed below
+
+    def encode_prompt(*args, **kwargs):
+        pipe.text_encoder = pipe.text_encoder.to("cuda", non_blocking=True)
+        try:
+            return inner(*args, **kwargs)
+        finally:
+            # finally, not just on success: leaving 11 GB stranded after a failed
+            # encode would make the retry the thing that OOMs.
+            pipe.text_encoder = pipe.text_encoder.to("cpu")
+            gc.collect()
+            torch.cuda.empty_cache()
+
+    # Set on the INSTANCE: all four generate_* methods call self.encode_prompt.
+    pipe.encode_prompt = encode_prompt
+    log("offload: the DiT and VAE are resident; the text encoder (~11 GB) comes over "
+        "only to encode the prompt")
+    return pipe
+
+
 def _load_quantized_base_dit(root, subfolder, **kwargs):
     """Load an INT8-quantised BASE DiT without ever holding the bf16 model.