longcat: the INT8 build must not take the host down (v0.2.65)

Run against the real weights for the first time, it exhausted all 54 GB of RAM
and all 4 GB of swap. MemAvailable reached 0 and I killed it to keep the server
alive — the production container was two minutes from being the thing the OOM
killer chose.

It had loaded all six fp32 shards (13m42s) and reached "quantising every Linear to
INT8". The peak is not the bf16 model alone: the INT8 copies are built while the
originals are still referenced, and save_quantized_state_dict then materialises a
full state dict on top. Roughly 2x bf16, not 1x.

Three changes:

- It now REFUSES when the host cannot take it, with the numbers: "~42 GB needed,
  34 GB available". Estimated from the shards on disk (fp32, so bf16 is half)
  times 1.5, plus 4 GB for the interpreter and torch. min_free_gb overrides it for
  someone who knows better. A job that can kill the server should not be reachable
  from a button without this.
- The load streams through a meta-device skeleton (low_cpu_mem_usage=True) instead
  of building a second full copy while loading.
- gc.collect() between the phases, and the model is dropped before the save
  returns — the replaced bf16 weights are ~13 GB of garbage at exactly the moment
  the save wants to allocate.

What this means on THIS host, honestly: 42 GB needed against 54 GB total with the
server using ~20 GB. It will now refuse rather than try. Quantising here needs the
server stopped, or a lower-peak implementation that never materialises the whole
model (quantise shard by shard, write as it goes) — which is the real fix and is
not in this commit.

22 tests.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0182DhR1eFNPQyGmDrpHbedY

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

tests/test_longcat_quantize.py

@@ -142,3 +142,58 @@ def test_bf16_is_never_gated(tmp_path):
 def test_the_service_passes_the_checkpoint_to_the_gate():
     src = SERVICE.read_text(encoding="utf-8")
     assert "checkpoint_dir=checkpoint" in src, "the gate cannot see the built weights"
+
+
+# ── it must not take the machine down ───────────────────────────────────────
+# Measured: on a 54 GB host this filled RAM and all 4 GB of swap and had to be
+# killed to keep the server alive. The peak holds the bf16 model, the INT8 copy
+# and the state dict the save builds.
+
+import importlib.util as _ilu2
+
+_qspec = _ilu2.spec_from_file_location("lc_quant", SCRIPT)
+LQ = _ilu2.module_from_spec(_qspec)
+_qspec.loader.exec_module(LQ)
+
+
+def test_it_estimates_the_peak_from_the_weights_on_disk(tmp_path):
+    d = tmp_path / "dit"
+    d.mkdir()
+    # 8 GB of fp32 shards -> 4 GB bf16 -> 1.5x + 4 floor
+    (d / "a.safetensors").write_bytes(b"\0" * 1024)
+    size = 8 * 1024 ** 3
+    import os as _os
+    _os.truncate(d / "a.safetensors", size)
+    got = LQ._needed_gb(str(tmp_path), "dit")
+    assert abs(got - (4 * 1.5 + 4)) < 0.1, got
+
+
+def test_an_unreadable_checkpoint_does_not_block_the_run(tmp_path):
+    """0 means 'unknown', and unknown must not refuse."""
+    assert LQ._needed_gb(str(tmp_path), "nope") == 0.0
+
+
+def test_available_memory_is_read_from_the_kernel():
+    got = LQ._available_gb()
+    assert got >= 0.0
+    if got:
+        assert got < 10000, "that is not gigabytes"
+
+
+def test_the_job_refuses_rather_than_exhausting_the_host():
+    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 "min_free_gb" in src
+
+
+def test_the_load_streams_instead_of_building_a_second_copy():
+    src = SCRIPT.read_text(encoding="utf-8")
+    assert "low_cpu_mem_usage=True" in src
+
+
+def test_memory_is_released_between_the_phases():
+    src = SCRIPT.read_text(encoding="utf-8")
+    assert src.count("gc.collect()") >= 3
+    assert "del dit" in src

tools/longcat_quantize.py

@@ -20,6 +20,7 @@ Invoked by codai/api/longcat_worker.py, never by hand:
 """
 
 import argparse
+import gc
 import json
 import os
 import sys
@@ -27,6 +28,35 @@ import time
 import traceback
 
 
+def _available_gb() -> float:
+    """Host memory actually available, or 0.0 if it cannot be read."""
+    try:
+        with open("/proc/meminfo", encoding="utf-8") as fh:
+            for line in fh:
+                if line.startswith("MemAvailable:"):
+                    return int(line.split()[1]) / (1024 * 1024)
+    except OSError:
+        pass
+    return 0.0
+
+
+def _needed_gb(source: str, subfolder: str) -> float:
+    """A floor for the peak, from the weights on disk.
+
+    The shards are fp32, so the bf16 model is about half of them; the peak then
+    holds that plus the INT8 copy plus the state dict the save builds. 1.5x the
+    bf16 size, with a 4 GB floor for the interpreter and torch itself.
+    """
+    d = os.path.join(source, subfolder)
+    try:
+        on_disk = sum(os.path.getsize(os.path.join(d, f)) for f in os.listdir(d)
+                      if f.endswith(".safetensors"))
+    except OSError:
+        return 0.0
+    bf16 = on_disk / 2 / (1024 ** 3)
+    return bf16 * 1.5 + 4
+
+
 def main(argv=None) -> int:
     ap = argparse.ArgumentParser(description="Quantise a LongCat DiT to INT8")
     ap.add_argument("--job", required=True, help="path to the job JSON")
@@ -64,19 +94,43 @@ def main(argv=None) -> int:
         if os.path.isdir(out_dir) and not job.get("overwrite"):
             return finish(True, path=out_dir, skipped="already quantised")
 
-        emit(step=1, total=4, message=f"loading the bf16 DiT from {subfolder}/")
-        # On CPU on purpose: the bf16 DiT is ~27 GB and the point of this job is to
-        # make it FIT the card, so it must not need the card to get there.
+        # Refuse rather than take the machine down. Measured the hard way: on a 54 GB
+        # host this filled RAM AND all 4 GB of swap, and had to be killed to keep the
+        # server alive. The peak is roughly the bf16 model plus the INT8 copy plus a
+        # full state dict at save time, so ~2x the bf16 size is the honest floor.
+        need_gb = float(job.get("min_free_gb") or 0) or _needed_gb(source, subfolder)
+        free_gb = _available_gb()
+        if free_gb and free_gb < need_gb:
+            return finish(False, error=(
+                f"not enough host memory to quantise: ~{need_gb:.0f} GB needed, "
+                f"{free_gb:.0f} GB available. The bf16 model, the INT8 copy and the "
+                f"state dict are all resident at the peak. Free memory (stop other "
+                f"models) or raise min_free_gb if you know better."))
+
+        emit(step=1, total=4,
+             message=f"loading the bf16 DiT from {subfolder}/ (~{need_gb:.0f} GB needed)")
+        # On CPU on purpose: the point of this job is to make a model that FITS the
+        # card, so getting there must not need the card. low_cpu_mem_usage streams the
+        # shards through a meta-device skeleton instead of building a second full copy
+        # while loading.
         dit = LongCatVideoTransformer3DModel.from_pretrained(
-            source, subfolder=subfolder, torch_dtype=torch.bfloat16)
+            source, subfolder=subfolder, torch_dtype=torch.bfloat16,
+            low_cpu_mem_usage=True)
         dit = dit.to("cpu").eval()
+        gc.collect()
 
         emit(step=2, total=4, message="quantising every Linear to INT8")
+        with torch.no_grad():
             dit = quantize_model(dit)
+        # The bf16 weights the QuantizedLinear layers replaced are garbage now; they
+        # are ~13 GB and the save below allocates again, so collect before it.
+        gc.collect()
 
         emit(step=3, total=4, message=f"writing {out_name}/")
         save_quantized_state_dict(
             dit, out_dir, config_source_dir=os.path.join(source, subfolder))
+        del dit
+        gc.collect()
 
         size = sum(
             os.path.getsize(os.path.join(out_dir, f))