longcat: quantise the text encoder, so nothing has to leave the card (v0.2.77)

The DiT has upstream's INT8 path. The UMT5-XXL text encoder does not, and at
bf16 it is ~11 GB — 45% of why a 13.6 GB INT8 DiT still overflows a 24 GB card
(13.6 + 11.2 + VAE = 25.1 GB before a single activation). Offload worked around
it by making the DiT and the encoder take turns; quantising the encoder removes
the reason to.

Measured against the real checkpoint, what the card must hold:
  bf16 encoder, offloaded   13.9 GB
  bf16 encoder, resident    25.2 GB   <- the configuration that OOMs
  int8 encoder, resident    19.5 GB
  nf4  encoder, resident    16.7 GB

So `text_encoder_quant` is a setting: none | int8 | nf4 | fp4, per model or
server-wide, validated at startup rather than after a seven-minute load.
bitsandbytes is pinned in requirements-longcat.txt and asserted by both the pod
image build and the local image's smoke test.

A quantised encoder is RESIDENT and never offloaded: bitsandbytes quantises as
the weights land on CUDA and is not built to shuttle them back — and at 2.8 GB
it does not need to. That is detected from the MODEL, not the config, so a
pipeline loaded quantised cannot be offloaded by a config that changed since,
and defensively, so a component that cannot be walked is simply not quantised.

Inserting these helpers by anchor dropped one of them between
@contextlib.contextmanager and the function it was meant to decorate, so
_encoder_is_quantised returned a context manager — truthy, making every pipeline
look quantised — and _bsa_for quietly stopped being one. Both now have a test
asserting their shape, because the symptom was five unrelated failures.

Unverified: bitsandbytes has not executed on this GPU. A throwaway container
cannot reach the card here (torch reports no CUDA devices under both --gpus all
and --runtime=nvidia, though the production container's own CUDA is fine), so
the NF4 path is proven only by its unit tests until a service run exercises it.
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.76"
+__version__ = "0.2.77"
 __tagline__ = "Complete orchestration, distribution and escalation of remotizable advanced inference."
 
 # Configure the CUDA caching allocator BEFORE torch is imported anywhere.

codai/api/longcat_worker.py

@@ -46,6 +46,7 @@ Per-model config keys (models.json entry, all optional)::
     "offload_strategy": "",        # '' | model | sequential
     "bsa": "off",                  # block-sparse attention: on | off | auto
     "base_model": "",              # avatar families borrow tokenizer/text_encoder/vae
+    "text_encoder_quant": "none",   # none | int8 | nf4 | fp4 (keeps it resident)
     "cp_split_hw": "",             # context-parallel tile, e.g. "1x2"
     "used_vram_gb": 0,             # what to reserve before starting; 0 = measured
     "gpu_device": 0
@@ -500,6 +501,13 @@ def ensure_service(model_path: str, config: dict = None,
         base = str(config.get("base_model") or "").strip()
         if base:
             cmd += ["--base-model", base]
+        # Quantising the text encoder is what lets the whole pipeline stay resident on
+        # a 24 GB card; per model, because it trades prompt fidelity for the room.
+        te_q = str(config.get("text_encoder_quant") or
+                   (getattr(sec, "text_encoder_quant", "") if sec is not None else "")
+                   or "").strip()
+        if te_q:
+            cmd += ["--text-encoder-quant", te_q]
         if sec is not None and str(getattr(sec, "extra_args", "") or "").strip():
             import shlex
             cmd += shlex.split(str(sec.extra_args))

codai/models/manager.py

@@ -4387,6 +4387,20 @@ class MultiModelManager:
         vae = _resident_gb("vae")
         encoder = _resident_gb("text_encoder")
 
+        # A quantised text encoder is RESIDENT and much smaller, so it neither takes
+        # turns with the DiT nor costs its bf16 size.
+        te_quant = str(cfg.get("text_encoder_quant") or "none").strip().lower()
+        try:
+            from codai.api.longcat_worker import common as _lc_common
+            te_bpe = _lc_common().component_quant_bpe(te_quant)
+        except Exception:
+            te_bpe = None
+        if te_bpe is not None:
+            stored = self._safetensors_storage_bpe(os.path.join(root, "text_encoder"))
+            if stored > 0:
+                encoder = _resident_gb("text_encoder") * (te_bpe / float(load_bpe))
+            return dit + encoder + vae
+
         offload = str(cfg.get("offload_strategy") or "").strip().lower()
         if offload and offload != "none":
             return max(dit, encoder) + vae

packaging/linux/smoke_test_services.sh

@@ -83,10 +83,10 @@ if "${DK[@]}" exec "$NAME" sh -lc "test -d /opt/coderai/longcat_venv"; then
   # re-pointed at the wrong interpreter would pass a bare "import torch" while being
   # unusable for LongCat.
   if "${DK[@]}" exec "$NAME" /opt/coderai/longcat_venv/bin/python -c \
-      "import sys,torch,transformers,accelerate; assert sys.version_info[:2]==(3,10), sys.version; assert torch.__version__.startswith('2.6'), torch.__version__; assert transformers.__version__.startswith('4.41'), transformers.__version__" >/dev/null 2>&1; then
-    ok "longcat venv imports py3.10 + torch 2.6 + transformers 4.41 + accelerate"
+      "import sys,torch,transformers,accelerate,bitsandbytes; assert sys.version_info[:2]==(3,10), sys.version; assert torch.__version__.startswith('2.6'), torch.__version__; assert transformers.__version__.startswith('4.41'), transformers.__version__" >/dev/null 2>&1; then
+    ok "longcat venv imports py3.10 + torch 2.6 + transformers 4.41 + accelerate + bitsandbytes"
   else
-    bad "longcat venv" "py3.10/torch-2.6/transformers-4.41/accelerate check failed"
+    bad "longcat venv" "py3.10/torch-2.6/transformers-4.41/accelerate/bitsandbytes check failed"
   fi
 else
   echo "   (skipped: no longcat venv in this image)"

packaging/runpod/Dockerfile.capability-video-longcat

@@ -17,7 +17,7 @@ RUN python -m venv /opt/coderai/longcat_venv \
       assert sys.version_info[:2]==(3,10), sys.version; \
       assert torch.__version__.startswith('2.6'), torch.__version__; \
       assert transformers.__version__.startswith('4.41'), transformers.__version__; \
-      import accelerate" \
+      import accelerate, bitsandbytes" \
  && rm -rf /root/.cache/pip /tmp/requirements-longcat.txt
 
 # ---------------------------------------------------------------- runtime

requirements-longcat.txt

@@ -50,6 +50,13 @@ diffusers==0.35.1
 # and filled in place: the 5-shard load went from ~4m50s to seconds of it, and the host
 # RAM spike goes away. It is also what diffusers wants for any device_map work later.
 accelerate==1.1.1
+
+# Weight quantisation for the COMPONENTS, driven from the model config
+# (text_encoder_quant / vae_quant). The DiT has upstream's own INT8 path; the
+# UMT5-XXL text encoder does not, and at bf16 it is ~11 GB — 45% of why an INT8
+# DiT still does not fit a 24 GB card. NF4 puts it at ~2.8 GB. CUDA-only, which is
+# fine: these components only ever run on a card.
+bitsandbytes==0.45.5
 einops==0.8.0
 ftfy==6.2.0
 loguru==0.7.2

tests/test_longcat_component_quant.py

+"""Quantising the components that upstream does not quantise for us.
+
+The DiT has upstream's own INT8 path. The UMT5-XXL text encoder does not, and at bf16
+it is ~11 GB — 45% of why a 13.6 GB INT8 DiT still overflows a 24 GB card (13.6 + 11.2
++ VAE = 25.1 GB before activations). Quantised it is ~5.6 GB at INT8 or ~2.8 GB at NF4,
+which is the difference between offloading and not.
+"""
+import importlib.util
+import pathlib
+import sys
+
+import pytest
+
+ROOT = pathlib.Path(__file__).resolve().parents[1]
+COMMON = ROOT / "tools" / "longcat_common.py"
+SERVICE = ROOT / "tools" / "longcat_service.py"
+WORKER = ROOT / "codai" / "api" / "longcat_worker.py"
+
+
+@pytest.fixture
+def LC():
+    spec = importlib.util.spec_from_file_location("lc_cq", COMMON)
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod
+
+
+# ------------------------------------------------------------------ the contract
+
+def test_none_means_keep_the_pipeline_dtype(LC):
+    assert LC.component_quant_bpe("none") is None
+    assert LC.component_quant_bpe("") is None
+    assert LC.component_quant_bpe(None) is None
+
+
+@pytest.mark.parametrize("kind,bpe", [("int8", 1.0), ("nf4", 0.5), ("fp4", 0.5)])
+def test_the_widths_are_what_the_card_will_hold(LC, kind, bpe):
+    assert LC.component_quant_bpe(kind) == bpe
+
+
+def test_an_unknown_setting_is_refused_with_the_options(LC):
+    problems = LC.component_quant_problems("text_encoder_quant", "4bit-ish")
+    assert problems and "int8" in problems[0] and "nf4" in problems[0]
+
+
+def test_a_valid_setting_passes(LC):
+    for kind in ("none", "int8", "nf4", "fp4", "NF4", " int8 "):
+        assert LC.component_quant_problems("text_encoder_quant", kind) == []
+
+
+# ------------------------------------------------------------------ the arithmetic
+
+def test_quantising_the_encoder_is_what_makes_it_fit(LC):
+    """The whole justification, in numbers: 22 GB of fp32 on disk becomes 11 GB at bf16
+    and does not fit beside the DiT; at NF4 it does."""
+    disk_gb, load_bpe, stored_bpe = 22.0, 2.0, 4.0
+    bf16 = disk_gb * (load_bpe / stored_bpe)
+    nf4 = disk_gb * (LC.component_quant_bpe("nf4") / stored_bpe)
+    int8 = disk_gb * (LC.component_quant_bpe("int8") / stored_bpe)
+    dit, vae = 13.6, 0.25
+    assert dit + bf16 + vae > 23.5, "this is the configuration that OOMs"
+    assert dit + int8 + vae < 23.5
+    assert dit + nf4 + vae < 20
+
+
+# ------------------------------------------------------------------ the service
+
+def test_the_service_refuses_a_bad_value_before_loading_anything():
+    src = SERVICE.read_text(encoding="utf-8")
+    assert "component_quant_problems" in src
+    main = src[src.index('ap.add_argument("--text-encoder-quant"'):]
+    assert "SystemExit" in main, "a typo must fail at startup, not after a 7-minute load"
+
+
+def test_a_quantised_encoder_is_never_offloaded():
+    """bitsandbytes quantises as weights land on CUDA and is not built to move them
+    back; it is also small enough that nothing needs to take turns."""
+    src = SERVICE.read_text(encoding="utf-8")
+    body = src[src.index("def _place_pipeline"):src.index("def _load_quantized_base_dit")]
+    assert "_encoder_is_quantised(pipe)" in body
+    # it must return before reaching the swap wrapper
+    assert body.index("_encoder_is_quantised(pipe)") < body.index("swap=True")
+
+
+def test_quantisation_is_detected_from_the_model_not_the_config():
+    """A pipeline loaded quantised must not be offloaded by a config that changed since."""
+    src = SERVICE.read_text(encoding="utf-8")
+    body = src[src.index("def _encoder_is_quantised"):src.index("def _quant_config")]
+    assert "Linear4bit" in body and "Linear8bitLt" in body
+    assert "text_encoder_quant" not in body
+
+
+def test_missing_bitsandbytes_is_a_clear_error():
+    src = SERVICE.read_text(encoding="utf-8")
+    body = src[src.index("def _quant_config"):src.index("def _wrap_encode_prompt")]
+    assert "needs bitsandbytes" in body
+    assert "requirements-longcat.txt" in body, "say where to fix it"
+
+
+def test_a_quantised_load_passes_a_device_map():
+    """bitsandbytes quantises on the way to the card; without device_map the weights
+    never get there and the config is silently ignored."""
+    src = SERVICE.read_text(encoding="utf-8")
+    assert "device_map" in src[src.index("te_quant ="):src.index("UMT5EncoderModel.from_pretrained")+200]
+
+
+# ------------------------------------------------------------------ plumbing
+
+def test_it_is_configurable_per_model_and_server_wide():
+    src = WORKER.read_text(encoding="utf-8")
+    assert '"--text-encoder-quant"' in src
+    body = src[src.index('te_q = str(config.get("text_encoder_quant")'):]
+    # per-model first: one checkpoint trading prompt fidelity for room must not force
+    # that trade on every other model on the box
+    assert body.index('config.get("text_encoder_quant")') < \
+        body.index('getattr(sec, "text_encoder_quant"')
+
+
+def test_it_is_documented_in_the_config_keys():
+    src = WORKER.read_text(encoding="utf-8")
+    assert "text_encoder_quant" in src[:src.index("import collections")]
+
+
+def test_bitsandbytes_is_pinned_and_verified_in_both_images():
+    reqs = (ROOT / "requirements-longcat.txt").read_text(encoding="utf-8")
+    assert "bitsandbytes==" in reqs
+    pod = (ROOT / "packaging/runpod/Dockerfile.capability-video-longcat").read_text()
+    assert "bitsandbytes" in pod
+    smoke = (ROOT / "packaging/linux/smoke_test_services.sh").read_text()
+    assert "bitsandbytes" in smoke
+
+
+# ------------------------------------------------------------------ shape of the helpers
+# Inserting these functions by anchor put one of them between @contextlib.contextmanager
+# and the function it was meant to decorate, so _encoder_is_quantised returned a context
+# manager (truthy — every pipeline looked quantised) and _bsa_for stopped being one.
+# Cheap guards, because the symptom was five unrelated test failures.
+
+def test_the_quantisation_check_returns_a_bool(service_mod):
+    import inspect
+    fn = service_mod._encoder_is_quantised
+    assert not hasattr(fn, "__wrapped__"), "something is decorating it"
+    assert not inspect.isgeneratorfunction(fn)
+
+    class _NoModules:
+        pass
+
+    class _Pipe:
+        text_encoder = _NoModules()
+
+    assert service_mod._encoder_is_quantised(_Pipe()) is False
+    assert service_mod._encoder_is_quantised(type("P", (), {"text_encoder": None})()) is False
+
+
+def test_the_bsa_guard_is_a_context_manager(service_mod):
+    import contextlib
+    cm = service_mod._bsa_for(type("P", (), {"dit": None})(), {})
+    assert hasattr(cm, "__enter__") and hasattr(cm, "__exit__")
+    with cm as active:
+        assert active is False          # no dit, so nothing to tile
+
+
+@pytest.fixture
+def service_mod(monkeypatch):
+    spec = importlib.util.spec_from_file_location("lc_svc_shape", SERVICE)
+    mod = importlib.util.module_from_spec(spec)
+    monkeypatch.setitem(sys.modules, spec.name, mod)
+    spec.loader.exec_module(mod)
+    return mod

tests/test_longcat_placement.py

@@ -23,16 +23,21 @@ SERVICE = ROOT / "tools" / "longcat_service.py"
 class FakeModule:
     """Records every device it is moved to, in order."""
 
-    def __init__(self, name):
+    def __init__(self, name, children=()):
         self.name = name
         self.device = "cpu"
         self.history = ["cpu"]
+        self.children = list(children) or [self]
 
     def to(self, device, non_blocking=False):
         self.device = str(device)
         self.history.append(str(device))
         return self
 
+    def modules(self):
+        """nn.Module's walk, which _encoder_is_quantised uses to spot bnb layers."""
+        return iter(self.children)
+
 
 class FakePipeline:
     """Upstream's shape: a plain object with its own to(), no offload hooks at all."""

tools/longcat_common.py

@@ -332,6 +332,31 @@ def bsa_problems(num_frames: int, height: int, width: int, chunk=None) -> list:
     return problems
 
 
+# Weight quantisation for the pipeline COMPONENTS, as bytes per element once loaded.
+# The DiT has upstream's own INT8 path (base_model_int8/); these cover the components
+# that do not, the UMT5-XXL text encoder above all: at bf16 it is ~11 GB, which is 45%
+# of why a 13.6 GB INT8 DiT still does not fit a 24 GB card.
+COMPONENT_QUANTS = {
+    "none": None,       # load at the pipeline dtype
+    "int8": 1.0,
+    "nf4": 0.5,         # 4 bits plus absmax/quant-map overhead, which rounds to this
+    "fp4": 0.5,
+}
+
+
+def component_quant_problems(name: str, value: str) -> list:
+    """Reasons this component cannot be quantised that way, or []."""
+    v = str(value or "none").strip().lower()
+    if v not in COMPONENT_QUANTS:
+        return [f"{name} must be one of {sorted(COMPONENT_QUANTS)}; got {value!r}"]
+    return []
+
+
+def component_quant_bpe(value: str):
+    """Bytes per element for a quantised component, or None to keep the load dtype."""
+    return COMPONENT_QUANTS.get(str(value or "none").strip().lower())
+
+
 def dit_subdir(family: str = "") -> str:
     """Which subfolder holds the bf16 transformer for this family."""
     return AVATAR_DIT_SUBDIR if family else BASE_DIT_SUBDIR

tools/longcat_service.py

@@ -62,6 +62,7 @@ _state = {
     "pipe": None, "model": "", "source": "", "dtype": "bfloat16",
     "family": "", "variant": "bf16", "use_int8": False, "use_distill": False,
     "cp_split_hw": [1, 1], "audio_encoder": None, "base_model": "",
+    "text_encoder_quant": "none",
 }
 _progress = {"active": False, "stage": "", "segment": 0, "segments": 0,
              "step": 0, "steps": 0}
@@ -116,8 +117,18 @@ def load_pipeline(checkpoint_dir: str, source_dir: str, dtype: str = "bfloat16",
             log(f"family '{family}': tokenizer/text_encoder/vae from {shared}")
 
         tokenizer = AutoTokenizer.from_pretrained(shared, subfolder="tokenizer")
+        te_quant = str(_state.get("text_encoder_quant") or "none").strip().lower()
+        te_kw = {"torch_dtype": td}
+        _te_cfg = _quant_config(te_quant, td)
+        if _te_cfg is not None:
+            # device_map is required for a quantised load: bitsandbytes quantises as the
+            # weights land on the card.
+            te_kw.update(quantization_config=_te_cfg, device_map={"": 0})
+            log(f"text encoder: {te_quant} "
+                f"(~{22.0 * LC.component_quant_bpe(te_quant) / 4:.1f} GB instead of "
+                f"~11 GB at bf16) — resident, never offloaded")
         text_encoder = UMT5EncoderModel.from_pretrained(
-            shared, subfolder="text_encoder", torch_dtype=td)
+            shared, subfolder="text_encoder", **te_kw)
         vae = AutoencoderKLWan.from_pretrained(shared, subfolder="vae", torch_dtype=td)
         scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             root, subfolder="scheduler")
@@ -322,7 +333,52 @@ def _attention_kwargs():
         "attention has no fallback path")
 
 
-@contextlib.contextmanager
+def _encoder_is_quantised(pipe) -> bool:
+    """Whether the text encoder carries bitsandbytes weights.
+
+    Asked of the MODEL rather than the config, so a pipeline that was loaded quantised
+    cannot be offloaded by a config that has since changed.
+    """
+    encoder = getattr(pipe, "text_encoder", None)
+    if encoder is None:
+        return False
+    if getattr(encoder, "is_quantized", False) or getattr(encoder, "is_loaded_in_4bit", False) \
+            or getattr(encoder, "is_loaded_in_8bit", False):
+        return True
+    walk = getattr(encoder, "modules", None)
+    if not callable(walk):
+        return False
+    return any(type(m).__name__ in ("Linear4bit", "Linear8bitLt", "Params4bit")
+               for m in walk())
+
+
+def _quant_config(kind: str, dtype):
+    """A BitsAndBytesConfig for a component, or None to load at the pipeline dtype.
+
+    Quantised weights live on the CARD: bitsandbytes 4-bit params are quantised on
+    their first move to CUDA and are not meant to be shuttled back to host RAM. That
+    is the point here rather than a limitation — a ~2.8 GB NF4 encoder fits beside a
+    13.6 GB INT8 DiT, so nothing has to take turns at all.
+    """
+    bpe = LC.component_quant_bpe(kind)
+    if bpe is None:
+        return None
+    try:
+        import bitsandbytes  # noqa: F401
+    except ImportError:
+        raise RuntimeError(
+            f"component quantisation {kind!r} needs bitsandbytes, which is not "
+            f"installed in this venv (it is pinned in requirements-longcat.txt)")
+    from transformers import BitsAndBytesConfig
+
+    kind = str(kind).strip().lower()
+    if kind == "int8":
+        return BitsAndBytesConfig(load_in_8bit=True)
+    return BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type=kind,
+                              bnb_4bit_compute_dtype=dtype,
+                              bnb_4bit_use_double_quant=True)
+
+
 def _wrap_encode_prompt(pipe, swap: bool):
     """Encode a prompt at most once per request, and (when offloading) swap for it.
 
@@ -397,6 +453,7 @@ def _wrap_encode_prompt(pipe, swap: bool):
     return pipe
 
 
+@contextlib.contextmanager
 def _bsa_for(pipe, ctx):
     """Run this pass with block-sparse attention only if the geometry allows it.
 
@@ -453,6 +510,23 @@ def _place_pipeline(pipe, offload):
         log("WARNING: no CUDA device visible — this will be extremely slow")
         return pipe
 
+    # A quantised encoder is already on the card and must stay there: bitsandbytes
+    # quantises as the weights land on CUDA and is not built to shuttle them back. It
+    # is also small enough not to need to — which is the whole point of quantising it.
+    if _encoder_is_quantised(pipe):
+        log("text encoder is quantised and resident; nothing takes turns on the card")
+        if pipe.dit is not None:
+            pipe.dit = pipe.dit.to("cuda", non_blocking=True)
+            if getattr(pipe.dit, "lora_dict", None):
+                for network in pipe.dit.lora_dict.values():
+                    for lora in network.loras:
+                        lora.to("cuda", non_blocking=True)
+        if pipe.vae is not None:
+            pipe.vae = pipe.vae.to("cuda", non_blocking=True)
+        pipe.device = "cuda"
+        _wrap_encode_prompt(pipe, swap=False)
+        return pipe
+
     if offload not in ("model", "sequential"):
         pipe.to("cuda")
         # Still worth caching: 12 segments would otherwise run 12 identical encodes.
@@ -905,6 +979,8 @@ def generate(body: dict) -> dict:
 
     if body.get("base_model"):
         _state["base_model"] = str(body["base_model"])
+    if body.get("text_encoder_quant"):
+        _state["text_encoder_quant"] = str(body["text_encoder_quant"])
     _state.update(variant=variant, use_int8=use_int8, use_distill=use_distill_lora)
     pipe = load_pipeline(checkpoint, body.get("source") or _state["source"],
                          body.get("dtype") or _state["dtype"],
@@ -1166,6 +1242,9 @@ def main(argv=None):
     ap.add_argument("--offload", default="", help="'' | model | sequential")
     ap.add_argument("--attention", default="xformers",
                     help="xformers | flash | flash3 (falls back to what is installed)")
+    ap.add_argument("--text-encoder-quant", default="none",
+                    help="none | int8 | nf4 | fp4 — quantise the UMT5-XXL text encoder "
+                         "so it can stay resident instead of being offloaded")
     ap.add_argument("--base-model", default="",
                     help="checkpoint holding tokenizer/text_encoder/vae for an avatar "
                          "family; blank uses the upstream base repo")
@@ -1191,6 +1270,11 @@ def main(argv=None):
     os.environ.setdefault("LONGCAT_ATTENTION", args.attention)
     os.environ.setdefault("LONGCAT_BSA", args.bsa)
     _state["base_model"] = args.base_model or ""
+    _problems = LC.component_quant_problems("text_encoder_quant",
+                                            args.text_encoder_quant)
+    if _problems:
+        raise SystemExit("; ".join(_problems))
+    _state["text_encoder_quant"] = args.text_encoder_quant or "none"
     # Validate both before anything slow happens: a bad value should not surface after
     # a multi-minute load.
     _bsa_enabled()