longcat: the quantised-encoder estimate was 3x too optimistic (v0.2.79)

bitsandbytes NF4 works — verified on the 3090, in the LongCat venv, loading the
real UMT5-XXL: 144 Linear4bit modules, is_loaded_in_4bit=True, finite embeddings
of the right shape. But it does not cost what v0.2.77 claimed.

bitsandbytes replaces Linear layers only. This encoder's 256k-token embedding
table and one linear per layer stay at the load dtype, leaving 6.13 GB that
cannot shrink. Measured on the card, against ~11.0 GB at bf16:

    nf4 + double quant   8.00 GB      (the bits-per-weight estimate said 2.8)
    int8                 9.76 GB      (it said 5.6)

So the estimate now comes from those measurements rather than the bit width,
because under-reserving by 3x is an OOM at load time. What the card must hold:

    bf16 encoder, offloaded   13.9 GB    9.7 GB headroom
    nf4  encoder, resident    22.1 GB    1.4 GB headroom
    int8 encoder, resident    24.0 GB    does not fit
    bf16 encoder, resident    25.2 GB    does not fit

Which inverts the recommendation this was built for: an NF4 encoder does fit a
24 GB card, but leaves less room for activations than a 29-frame clip already
used, so it would OOM on anything longer. Offload plus the per-request embedding
cache from v0.2.76 keeps 9.7 GB free and touches the encoder once, so it stays
the better configuration here. The setting remains, correctly sized, for cards
where it is the right trade.

The model page now states the measured figures too, since "nf4 — ~2.8 GB" would
have invited exactly the configuration that OOMs.
Co-Authored-By: 's avatarClaude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0185zEBiy3KWB37xNGSrxL7w
parent 43ce2008
......@@ -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.78"
__version__ = "0.2.79"
__tagline__ = "Complete orchestration, distribution and escalation of remotizable advanced inference."
# Configure the CUDA caching allocator BEFORE torch is imported anywhere.
......
......@@ -708,9 +708,9 @@ window.__DEFAULT_WHISPER_SERVER_PATH__ = {{ default_whisper_server_path|tojson }
<div class="form-row" style="margin:0"><label class="form-label">Text encoder</label>
<select id="cfg-longcat-te-quant" class="form-input">
<option value="">bf16 — ~11 GB, needs offload</option>
<option value="int8">int8 — ~5.6 GB, stays resident</option>
<option value="nf4">nf4 — ~2.8 GB, stays resident</option>
<option value="fp4">fp4 — ~2.8 GB, stays resident</option>
<option value="int8">int8 — ~9.8 GB, too big to stay resident</option>
<option value="nf4">nf4 — ~8.0 GB, resident but tight</option>
<option value="fp4">fp4 — ~8.0 GB, resident but tight</option>
</select></div>
<div class="form-row" style="margin:0"><label class="form-label">Block-sparse attention</label>
<select id="cfg-longcat-bsa" class="form-input">
......
......@@ -4392,13 +4392,14 @@ class MultiModelManager:
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)
te_fraction = _lc_common().component_quant_fraction(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))
te_fraction = None
if te_fraction is not None:
# A MEASURED fraction of the bf16 size, not the bit width: bitsandbytes
# quantises Linear layers only, and this encoder's embedding table is 6 GB
# that stays put. The bits-per-weight figure was 3x too optimistic.
encoder = encoder * float(te_fraction)
return dit + encoder + vae
offload = str(cfg.get("offload_strategy") or "").strip().lower()
......
......@@ -50,17 +50,39 @@ def test_a_valid_setting_passes(LC):
# ------------------------------------------------------------------ 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
def test_the_measured_fractions_are_used_not_the_bit_width(LC):
"""MEASURED on a 3090, because bitsandbytes replaces Linear layers only.
UMT5-XXL has 144 of them; its 256k-token embedding table and one linear per layer
stay at the load dtype, leaving 6.13 GB that cannot shrink. The bits-per-weight
arithmetic predicted 2.8 GB for NF4 and the card showed 8.0 — so the estimate has
to come from the measurement, or it under-reserves by 3x and OOMs.
"""
bf16 = 11.0
assert LC.component_quant_fraction("none") is None
nf4 = bf16 * LC.component_quant_fraction("nf4")
int8 = bf16 * LC.component_quant_fraction("int8")
assert 7.5 < nf4 < 8.5, nf4
assert 9.3 < int8 < 10.3, int8
# and the naive figure it replaced must not be what anyone reserves from: it is
# 2.75 GB against a measured 8.0, which is what would have under-reserved by ~3x
naive = 22.0 * LC.component_quant_bpe("nf4") / 4
assert nf4 > 2.5 * naive, f"measured {nf4:.1f} vs naive {naive:.1f}"
def test_quantising_the_encoder_does_not_make_the_pipeline_fit_comfortably(LC):
"""The conclusion the measurement forced: with a 13.6 GB INT8 DiT and a 0.25 GB VAE,
an NF4 encoder leaves ~1.4 GB for activations and an INT8 one does not fit at all.
Offload keeps 9.7 GB free, so it remains the better configuration on a 24 GB card.
"""
dit, vae, usable, bf16 = 13.6, 0.25, 23.56, 11.0
nf4_total = dit + bf16 * LC.component_quant_fraction("nf4") + vae
int8_total = dit + bf16 * LC.component_quant_fraction("int8") + vae
assert int8_total > usable, "int8 encoder resident does not fit"
assert nf4_total < usable, "nf4 encoder resident does fit"
assert usable - nf4_total < 2.0, "but with almost no room for activations"
offloaded = dit + vae
assert usable - offloaded > 8, "offload has the headroom"
# ------------------------------------------------------------------ the service
......
......@@ -343,6 +343,23 @@ COMPONENT_QUANTS = {
"fp4": 0.5,
}
# What a quantised component ACTUALLY costs, as a fraction of its bf16 size — MEASURED
# on UMT5-XXL on a 3090, not derived from the bits per weight. bitsandbytes replaces
# Linear layers only: 144 of them here, while the 256k-token embedding table and one
# linear per layer stay at the load dtype. That leaves 6.13 GB unquantisable, so the
# bits-per-weight arithmetic was ~3x too optimistic (it predicted 2.8 GB for NF4; the
# card sees 8.0).
COMPONENT_QUANT_RESIDENT_FRACTION = {
"int8": 0.89, # 9.76 GB of an 11.0 GB bf16 encoder
"nf4": 0.73, # 8.00 GB
"fp4": 0.73,
}
def component_quant_fraction(value: str):
"""Fraction of the bf16 size a quantised component keeps, or None for no change."""
return COMPONENT_QUANT_RESIDENT_FRACTION.get(str(value or "none").strip().lower())
def component_quant_problems(name: str, value: str) -> list:
"""Reasons this component cannot be quantised that way, or []."""
......
......@@ -125,8 +125,9 @@ def load_pipeline(checkpoint_dir: str, source_dir: str, dtype: str = "bfloat16",
# 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")
f"(~{11.0 * LC.component_quant_fraction(te_quant):.1f} GB instead of "
f"~11 GB at bf16 — the embedding table does not quantise) "
f"— resident, never offloaded")
text_encoder = UMT5EncoderModel.from_pretrained(
shared, subfolder="text_encoder", **te_kw)
vae = AutoencoderKLWan.from_pretrained(shared, subfolder="vae", torch_dtype=td)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment