longcat: INT8 weights that can be built AND loaded (v0.2.66)

Building the INT8 DiT needed ~42 GB because it materialised the bf16 model
first; loading it needed ~55 GB for the same reason, plus every shard
accumulated into one dict and then copied in. Both took the host down.

Both now stream. The quantiser reads tensors one at a time through safe_open
and flushes 4 GB shards, enumerating the Linear layers on a meta-device
skeleton instead of a real model. _load_quantized_base_dit builds its skeleton
on meta too, then assigns each shard into it as it is read.

Measured on the 13.6B DiT: 51 GB bf16 -> 13.6 GB INT8, ~4 GB peak to write,
and the load is 20 s at 0.6 GB RSS (safetensors mmap) with all 437 quantised
Linears present and nothing stranded on meta.

A checkpoint that leaves tensors on the meta device is now an error naming
them, rather than a model that fails later somewhere unrelated.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QXD9yKKFZHAB9t1LfWQWwz

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.65"
+__version__ = "0.2.66"
 __tagline__ = "Complete orchestration, distribution and escalation of remotizable advanced inference."
 
 # Configure the CUDA caching allocator BEFORE torch is imported anywhere.

tests/test_longcat_int8_load.py

+"""Loading the INT8 LongCat DiT must not cost what the bf16 model costs.
+
+The straightforward way to load a quantised checkpoint — instantiate the model,
+swap the Linears, read every shard into one dict, copy it in — needs the full
+bf16 model (~27 GB for this 13.6B DiT) plus the weights twice before it is done.
+On a 54 GB host that is an out-of-memory kill, which is exactly what happened.
+
+So the contract these tests hold is about HOW it loads, not only that it loads:
+the skeleton is built on the meta device and each shard is streamed in and
+assigned. A fake DiT records the device it was constructed on, so an
+implementation that allocates for real fails here instead of on the host.
+"""
+import importlib.util
+import json
+import pathlib
+import sys
+
+import pytest
+
+torch = pytest.importorskip("torch")
+import torch.nn as nn
+from safetensors.torch import save_file
+
+ROOT = pathlib.Path(__file__).resolve().parents[1]
+SERVICE = ROOT / "tools" / "longcat_service.py"
+
+HIDDEN = 8
+DEPTH = 2
+
+
+class FakeDiT(nn.Module):
+    """Shaped like the real thing where it matters: a skipped final_layer.linear,
+    Linears inside numbered blocks, and plain parameters that are not quantised."""
+
+    construction_devices: list = []
+
+    def __init__(self, hidden_size=HIDDEN, depth=DEPTH, **kwargs):
+        super().__init__()
+        FakeDiT.construction_devices.append(torch.empty(1).device.type)
+        self.kwargs_seen = kwargs
+        self.blocks = nn.ModuleList(
+            nn.ModuleDict({"attn": nn.Linear(hidden_size, hidden_size, bias=False)})
+            for _ in range(depth)
+        )
+        self.norm = nn.LayerNorm(hidden_size)
+        self.final_layer = nn.ModuleDict(
+            {"linear": nn.Linear(hidden_size, hidden_size, bias=True)}
+        )
+
+
+class FakeQuantizedLinear(nn.Module):
+    """Upstream's layout: int8 weights and a per-channel scale, both buffers."""
+
+    def __init__(self, in_features, out_features, bias=False):
+        super().__init__()
+        self.register_buffer("weight_int8",
+                             torch.zeros(out_features, in_features, dtype=torch.int8))
+        self.register_buffer("weight_scale",
+                             torch.zeros(out_features, dtype=torch.float32))
+        if bias:
+            self.register_buffer("bias", torch.zeros(out_features, dtype=torch.bfloat16))
+
+
+@pytest.fixture
+def service(monkeypatch):
+    """tools/longcat_service.py with a fake longcat_video package behind it."""
+    pkg = type(sys)("longcat_video")
+    pkg.__path__ = []
+    modules = type(sys)("longcat_video.modules")
+    modules.__path__ = []
+    quant = type(sys)("longcat_video.modules.quantization")
+    quant.QuantizedLinear = FakeQuantizedLinear
+    quant.DEFAULT_SKIP_PATTERNS = {"final_layer.linear"}
+    dit = type(sys)("longcat_video.modules.longcat_video_dit")
+    dit.LongCatVideoTransformer3DModel = FakeDiT
+    for name, mod in (("longcat_video", pkg),
+                      ("longcat_video.modules", modules),
+                      ("longcat_video.modules.quantization", quant),
+                      ("longcat_video.modules.longcat_video_dit", dit)):
+        monkeypatch.setitem(sys.modules, name, mod)
+
+    spec = importlib.util.spec_from_file_location("lc_service_under_test", SERVICE)
+    mod = importlib.util.module_from_spec(spec)
+    monkeypatch.setitem(sys.modules, spec.name, mod)
+    spec.loader.exec_module(mod)
+    FakeDiT.construction_devices.clear()
+    return mod
+
+
+def _expected_tensors():
+    """What the quantiser writes for a FakeDiT: int8 pairs for the two block
+    Linears, bf16 for the skipped final layer, plain params for the norm."""
+    out = {}
+    for i in range(DEPTH):
+        out[f"blocks.{i}.attn.weight_int8"] = torch.full(
+            (HIDDEN, HIDDEN), i + 1, dtype=torch.int8)
+        out[f"blocks.{i}.attn.weight_scale"] = torch.full(
+            (HIDDEN,), 0.5 * (i + 1), dtype=torch.float32)
+    out["norm.weight"] = torch.ones(HIDDEN)
+    out["norm.bias"] = torch.zeros(HIDDEN)
+    out["final_layer.linear.weight"] = torch.full((HIDDEN, HIDDEN), 3.0)
+    out["final_layer.linear.bias"] = torch.full((HIDDEN,), 4.0)
+    return out
+
+
+def _write_checkpoint(qdir, tensors, shards=2):
+    qdir.mkdir(parents=True, exist_ok=True)
+    (qdir / "config.json").write_text(json.dumps(
+        {"hidden_size": HIDDEN, "depth": DEPTH, "_class_name": "ignored"}))
+    keys = list(tensors)
+    weight_map = {}
+    per = (len(keys) + shards - 1) // shards
+    for n in range(shards):
+        chunk = keys[n * per:(n + 1) * per]
+        if not chunk:
+            continue
+        name = f"quantized_model-{n + 1:05d}-of-{shards:05d}.safetensors"
+        save_file({k: tensors[k].contiguous() for k in chunk}, str(qdir / name))
+        weight_map.update({k: name for k in chunk})
+    (qdir / "quantized_model.safetensors.index.json").write_text(
+        json.dumps({"metadata": {}, "weight_map": weight_map}))
+
+
+def test_the_skeleton_is_built_on_the_meta_device(service, tmp_path):
+    """The load that took the host down built the bf16 model for real first."""
+    _write_checkpoint(tmp_path / "base_model_int8", _expected_tensors())
+
+    service._load_quantized_base_dit(str(tmp_path), "base_model_int8")
+
+    assert FakeDiT.construction_devices == ["meta"], (
+        "the DiT was allocated for real before any INT8 weight was read — that is "
+        "~27 GB of bf16 on the real model")
+
+
+def test_every_tensor_ends_up_real_and_correct(service, tmp_path):
+    """A meta skeleton is only useful if the streamed shards actually land in it."""
+    tensors = _expected_tensors()
+    _write_checkpoint(tmp_path / "base_model_int8", tensors)
+
+    model = service._load_quantized_base_dit(str(tmp_path), "base_model_int8")
+
+    state = dict(model.named_parameters())
+    state.update(model.named_buffers())
+    assert not [n for n, t in state.items() if t.is_meta]
+    assert set(state) == set(tensors)
+    for name, want in tensors.items():
+        assert torch.equal(state[name], want), name
+
+
+def test_the_skipped_layer_keeps_its_bf16_linear(service, tmp_path):
+    """final_layer.linear is in DEFAULT_SKIP_PATTERNS: quantising it wrecks output."""
+    _write_checkpoint(tmp_path / "base_model_int8", _expected_tensors())
+
+    model = service._load_quantized_base_dit(str(tmp_path), "base_model_int8")
+
+    assert isinstance(model.final_layer["linear"], nn.Linear)
+    assert all(isinstance(b["attn"], FakeQuantizedLinear) for b in model.blocks)
+
+
+def test_a_checkpoint_missing_tensors_is_an_error_not_a_meta_model(service, tmp_path):
+    """Returning a model with tensors still on meta hands the caller something that
+    fails much later, somewhere unrelated. Say so here."""
+    tensors = _expected_tensors()
+    del tensors["norm.weight"]
+    _write_checkpoint(tmp_path / "base_model_int8", tensors)
+
+    with pytest.raises(RuntimeError, match="meta device"):
+        service._load_quantized_base_dit(str(tmp_path), "base_model_int8")
+
+
+def test_extra_tensors_are_tolerated(service, tmp_path):
+    """A newer quantiser writing something this config has no slot for should warn,
+    not abort a load that is otherwise complete."""
+    tensors = _expected_tensors()
+    tensors["blocks.0.attn.weight_absmax"] = torch.ones(HIDDEN)
+    _write_checkpoint(tmp_path / "base_model_int8", tensors)
+
+    model = service._load_quantized_base_dit(str(tmp_path), "base_model_int8")
+
+    assert not [n for n, t in model.named_buffers() if t.is_meta]
+
+
+def test_shards_are_streamed_one_at_a_time(service, tmp_path):
+    """Accumulating every shard into one dict before loading doubles the peak."""
+    src = SERVICE.read_text(encoding="utf-8")
+    body = src[src.index("def _load_quantized_base_dit"):
+               src.index("def _load_audio_encoder")]
+    assert "safe_open" in body, "whole-shard load_file holds a shard in memory twice"
+    assert "assign=True" in body, "without assign= the meta skeleton is copied into"
+    assert "load_file" not in body

tests/test_longcat_quantize.py

@@ -36,9 +36,11 @@ def test_it_never_replaces_the_bf16_weights():
 
 def test_it_quantises_on_the_cpu():
     """The whole point is to make a model that does not fit the card — so getting
-    there must not need the card."""
+    there must not need the card. The streaming path reads tensors straight to CPU
+    and never touches CUDA."""
     src = SCRIPT.read_text(encoding="utf-8")
-    assert 'to("cpu")' in src
+    assert 'device="cpu"' in src
+    assert "cuda" not in src.lower()
 
 
 def test_an_existing_build_is_not_silently_redone():
@@ -180,20 +182,73 @@ def test_available_memory_is_read_from_the_kernel():
         assert got < 10000, "that is not gigabytes"
 
 
-def test_the_job_refuses_rather_than_exhausting_the_host():
+def test_the_job_still_checks_the_host_has_room():
     src = SCRIPT.read_text(encoding="utf-8")
-    assert "not enough host memory to quantise" in src
-    assert "_available_gb()" in src and "_needed_gb(" in src
-    # and the operator can override it knowingly
+    assert "not enough host memory" in src
+    assert "_available_gb()" in src
     assert "min_free_gb" in src
 
 
-def test_the_load_streams_instead_of_building_a_second_copy():
+def test_nothing_loads_the_weights_into_a_model():
+    """Superseded by the streaming rewrite: there is no load left to make cheap."""
     src = SCRIPT.read_text(encoding="utf-8")
-    assert "low_cpu_mem_usage=True" in src
+    assert "low_cpu_mem_usage" not in src
 
 
-def test_memory_is_released_between_the_phases():
+def test_memory_is_released_as_it_goes():
     src = SCRIPT.read_text(encoding="utf-8")
-    assert src.count("gc.collect()") >= 3
-    assert "del dit" in src
+    assert src.count("gc.collect()") >= 2
+    assert "del tensor" in src
+
+
+# ── the streaming rewrite ───────────────────────────────────────────────────
+# Materialising the model needed ~42 GB and exhausted a 54 GB host. The peak is
+# now one tensor plus the 4 GB shard being accumulated.
+
+def test_the_model_is_never_materialised():
+    src = SCRIPT.read_text(encoding="utf-8")
+    # matched as CALLS, so prose explaining what upstream does does not trip it
+    import re as _re
+    for call in ("from_pretrained", "quantize_model", "save_quantized_state_dict"):
+        assert not _re.search(r"^\s*(?:\w+\s*=\s*)?" + call + r"\(", src, _re.M), call
+
+
+def test_the_layout_is_learned_on_the_meta_device():
+    """Free: a 13.6B transformer costs no memory there, and it gives the exact set
+    of modules upstream would quantise rather than a guess from tensor shapes."""
+    src = SCRIPT.read_text(encoding="utf-8")
+    assert 'torch.device("meta")' in src
+    assert "DEFAULT_SKIP_PATTERNS" in src, "upstream's skip list must still apply"
+
+
+def test_tensors_are_streamed_one_at_a_time():
+    src = SCRIPT.read_text(encoding="utf-8")
+    assert "safe_open(" in src and "get_tensor(" in src
+
+
+def test_the_emitted_keys_match_what_the_loader_expects():
+    """load_quantized_dit builds QuantizedLinear modules whose buffers are named
+    weight_int8 / weight_scale / bias — the output has to use those exact names."""
+    src = SCRIPT.read_text(encoding="utf-8")
+    assert '.weight_int8"' in src
+    assert '.weight_scale"' in src
+    assert "quantized_model.safetensors.index.json" in src
+    assert "quantization_config.json" in src
+
+
+def test_the_quantisation_matches_upstreams_arithmetic():
+    """Per-channel symmetric, /127, clamped — same as QuantizedLinear.from_linear."""
+    src = SCRIPT.read_text(encoding="utf-8")
+    assert "abs().amax(dim=1).clamp(min=1e-8) / 127.0" in src
+    assert "round().clamp(-128, 127)" in src
+
+
+def test_shards_are_capped_so_the_buffer_cannot_grow_unbounded():
+    src = SCRIPT.read_text(encoding="utf-8")
+    assert "4 * 1024 * 1024 * 1024" in src
+    assert "_flush()" in src
+
+
+def test_the_memory_floor_is_now_a_sanity_check_not_the_constraint():
+    src = SCRIPT.read_text(encoding="utf-8")
+    assert "or 8.0" in src, "the 42 GB ceiling should be gone"

tools/longcat_service.py

@@ -205,17 +205,30 @@ def load_pipeline(checkpoint_dir: str, source_dir: str, dtype: str = "bfloat16",
 
 
 def _load_quantized_base_dit(root, subfolder, **kwargs):
-    """Load an INT8-quantised BASE DiT.
-
-    The same procedure as upstream's ``load_quantized_dit`` — build the model from
-    the saved config, swap every non-skipped Linear for its QuantizedLinear, then
-    load the sharded state dict — but instantiating LongCatVideoTransformer3DModel
-    instead of the avatar class upstream hardcodes. Its quantisation helpers are
-    reused as-is so the weight layout cannot drift from what wrote it.
+    """Load an INT8-quantised BASE DiT without ever holding the bf16 model.
+
+    Same end state as upstream's ``load_quantized_dit`` — the model built from the
+    saved config, every non-skipped Linear swapped for its QuantizedLinear, the
+    sharded state dict loaded — but instantiating LongCatVideoTransformer3DModel
+    instead of the avatar class upstream hardcodes, and getting there without the
+    peak that makes the straightforward version unusable.
+
+    Upstream instantiates the model for real, accumulates every shard into one
+    dict, then copies that into the model: for a 13.6B DiT that is ~27 GB of bf16
+    Linears before a single int8 byte is read, plus the weights twice over. It
+    exhausts a 54 GB host. Here the skeleton is built on the meta device (no
+    allocation), and each shard is streamed in and assigned, so the peak is one
+    shard above the final INT8 size.
+
+    Its quantisation helpers are reused as-is so the weight layout cannot drift
+    from what wrote it.
     """
+    import gc
+    import itertools
     import json as _json
+    import torch
     import torch.nn as _nn
-    from safetensors.torch import load_file as _load_file
+    from safetensors import safe_open
     from longcat_video.modules.quantization import QuantizedLinear, DEFAULT_SKIP_PATTERNS
     from longcat_video.modules.longcat_video_dit import LongCatVideoTransformer3DModel
 
@@ -226,35 +239,58 @@ def _load_quantized_base_dit(root, subfolder, **kwargs):
         config.pop(drop, None)
     config.update({k: v for k, v in kwargs.items() if v is not None})
 
-    model = LongCatVideoTransformer3DModel(**config)
-
-    replace = {}
-    for name, module in model.named_modules():
-        if isinstance(module, _nn.Linear) and not any(
-                pat in name for pat in DEFAULT_SKIP_PATTERNS):
-            replace[name] = QuantizedLinear(module.in_features, module.out_features,
-                                            bias=module.bias is not None)
-    for name, ql in replace.items():
-        parts = name.split(".")
-        parent = model
-        for part in parts[:-1]:
-            parent = getattr(parent, part)
-        setattr(parent, parts[-1], ql)
+    # The skeleton costs nothing on meta: shapes and module structure, no storage.
+    with torch.device("meta"):
+        model = LongCatVideoTransformer3DModel(**config)
+        replace = {}
+        for name, module in model.named_modules():
+            if isinstance(module, _nn.Linear) and not any(
+                    pat in name for pat in DEFAULT_SKIP_PATTERNS):
+                replace[name] = QuantizedLinear(module.in_features, module.out_features,
+                                                bias=module.bias is not None)
+        for name, ql in replace.items():
+            parts = name.split(".")
+            parent = model
+            for part in parts[:-1]:
+                parent = getattr(parent, part)
+            setattr(parent, parts[-1], ql)
+    log(f"INT8 DiT: {len(replace)} quantised Linear layers")
 
     index_path = os.path.join(qdir, "quantized_model.safetensors.index.json")
-    state = {}
     if os.path.exists(index_path):
         with open(index_path, encoding="utf-8") as fh:
             index = _json.load(fh)
-        for shard in sorted(set(index["weight_map"].values())):
-            state.update(_load_file(os.path.join(qdir, shard)))
+        shards = sorted(set(index["weight_map"].values()))
     else:
-        for shard in sorted(f for f in os.listdir(qdir) if f.endswith(".safetensors")):
-            state.update(_load_file(os.path.join(qdir, shard)))
-    missing, unexpected = model.load_state_dict(state, strict=False)
-    if missing:
-        log(f"WARNING: {len(missing)} missing tensors in the INT8 DiT "
-            f"(first: {missing[:3]})")
+        shards = sorted(f for f in os.listdir(qdir) if f.endswith(".safetensors"))
+    if not shards:
+        raise RuntimeError(f"no INT8 shards in {qdir}")
+
+    known = set(model.state_dict().keys())
+    unexpected = []
+    for shard in shards:
+        with safe_open(os.path.join(qdir, shard), framework="pt", device="cpu") as fh:
+            part = {}
+            for key in fh.keys():
+                if key in known:
+                    part[key] = fh.get_tensor(key)
+                else:
+                    unexpected.append(key)
+        # assign=True hands the loaded tensor to the module instead of copying into
+        # storage the meta skeleton does not have.
+        model.load_state_dict(part, strict=False, assign=True)
+        del part
+        gc.collect()
+        log(f"INT8 DiT: loaded {shard}")
+
+    stranded = [name for name, t in itertools.chain(model.named_parameters(),
+                                                    model.named_buffers())
+                if t.is_meta]
+    if stranded:
+        raise RuntimeError(
+            f"{len(stranded)} tensors were not in the INT8 checkpoint and are still "
+            f"on the meta device (first: {stranded[:5]}) — the weights in {qdir} do "
+            f"not match this model config")
     if unexpected:
         log(f"WARNING: {len(unexpected)} unexpected tensors in the INT8 DiT "
             f"(first: {unexpected[:3]})")