Laguna-S-2.1-W4A16 / LAGUNA-W4A16-BLOG.md
Relativ3pa1n's picture
Upload LAGUNA-W4A16-BLOG.md with huggingface_hub
183370c verified
|
Raw
History Blame Contribute Delete
8.68 kB

Laguna S 2.1 W4A16 on 8×3090 — quantizing a 256-expert MoE on a RAM-starved box

We quantized Laguna S 2.1 (poolside, 118B total / 8B activated, 256-expert MoE) to a clean W4A16 that serves on 8× RTX 3090 (192GB VRAM) at ~92 tok/s, 256K context, 7.3× concurrency via vLLM + Marlin — with thinking on, producing clean technical/markdown output. No official W4A16 for Ampere exists (poolside ships QAT-INT4 that loads poorly on Ampere, and FP8/NVFP4 which are Blackwell-only). This fills the gap for the 3090 crowd (tested on 8× RTX 3090; likely fine on other Ampere 24 GB+ cards, but untested).

It took four walls to get here. Each one looked like "the quant is broken." None of them were the final answer. This is the whole path, because the interesting parts are the dead ends.

Wall 1: llmcompressor's oneshot OOMs on a 256-expert MoE

llmcompressor.oneshot() — and every documented variant (load_quantizable_moe, load_context, sequential pipeline + disk offload) — runs linearize_moe on every MoE expert module in one pre-pipeline pass. Laguna has 256 routed experts across 47 routed-MoE layers. That global linearization transient blows past 400GB of host RAM at layer 4 of 47. Silent OOM-kill, no traceback. Five attempts, same wall, same line: Linearizing experts: 9% | 4/47.

Known class of problem (llm-compressor #1718, #1409). There is no --skip-linearize flag. The doc mitigations are AWQ-calibration-specific and don't apply to data-free RTN. pipeline="sequential" doesn't help because the linearize runs before the sequential pipeline.

Wall 2 → fix: quantize one MoE layer at a time (the restricted-RAM trick)

The script (w4a16-quant-laguna-s-2.1-perlayer.py) deliberately does not call oneshot(), load_quantizable_moe(), or linearize_moe(model). It duplicates llmcompressor.modeling.moe.linearize's core operation one MoE layer at a time:

linear_experts_cls = LinearExperts2D.get_linear_experts_cls(experts.__class__)
linear_moe = linear_experts_cls.from_experts_module(experts, config)
parent_layer.set_submodule(local_name, linear_moe)

Then it applies data-free RTN to the Linear modules inside that one linearized routed-experts block and immediately moves on. Peak host memory is about one layer's experts (5GB), not all 47 layers' (400GB+ transient). The BF16 base (219GB) loads via compressed_tensors.offload.load_offloaded_model (disk offload), so it never fully sits in RAM either.

What stays BF16 (ignored by construction): layer 0, all self_attn, the MoE router/gate, the shared expert, embed_tokens, lm_head. Only the 256 routed experts get quantized — the bulk of the params, the part that matters for size.

Output saves via llmcompressor's compressed-tensors save_pretrained, so config.json gets a quantization_config and vLLM auto-detects W4A16/Marlin.

That got a W4A16 out the door on a RAM-starved box. It did not get a clean one.

Wall 3: the output was garbage — and we blamed the quant (wrongly)

The first W4A16 served, but emitted severe number-soup: VG@141425700784791560079015245496..., bank996994878462523..., <|fim_pad|>. We called it RTN corruption and almost moved to W8A16.

The real cause: the BF16 base was a stale revision. poolside had re-pushed corrected weights to poolside/Laguna-S-2.1 on Jul 23–24 (Hugging Face lastModified was today); our download predated the fix. The community FP8 thread's fix — "pruned HF cache + re-downloaded the updated model → no issues" — was exactly this. One download-model.sh re-pull of the latest revision and the severe garbage vanished.

Lesson: severe garbage tokens (not mildly-off text) = corrupted/stale weights or a tokenizer mismatch, not aggressive quantization. Verify the base revision before blaming the quant.

Wall 4 → the real fix: group_size 128 was too coarse; 32 is the answer

The re-quantized W4A16 was mostly clean — short prompts perfect — but on longer, technical, markdown-heavy generation it still spliced stray number-tokens (Core Concept9, Queries (Q9:, DBC9FC10) and leaked special tokens (<|reserved_token_166473|>). Sampler flags didn't fix it; repetition_penalty 1.2 actually made it worse.

Root cause: group_size=128 (carried over from the Qwen3.5 recipe) is too coarse for Laguna's 256-expert MoE. Each group of 128 expert weights shared a single scale, so the per-group quant error was large enough to flip logits at token boundaries — and markdown/technical content hits those boundaries constantly. Creative prose stayed clean because it rarely does.

poolside's own reference INT4 uses group_size=32. We re-quantized at g32 (4× finer scales) and the number-tokens disappeared — clean markdown, correct math, thinking on. (API gotcha: QuantizationModifier's scheme field only accepts preset names, not a structured scheme dict — group_size is set via config_groups={"group_0": QuantizationScheme(targets=["Linear"], weights=preset_name_to_scheme("W4A16_ASYM",["Linear"]).weights.model_copy(update={"group_size":32}))}.)

The quant was never broken. The group_size was wrong for this architecture.

Wall 5 → throughput: --enforce-eager was leaving 12× on the table

With a clean model, decode was ~7 tok/s single-stream. On 8×3090. For an 8B-active model. That's absurdly slow.

Cause: the recipe ran --enforce-eager (no CUDA graphs). For a small-active MoE, the per-step Python + expert-routing-dispatch overhead dominates the tiny 8B compute — and eager mode pays that overhead every token. Dropping --enforce-eager (cudagraph capture succeeded at 0.90 util, no OOM) took decode to ~92 tok/s — a 12.5× jump. KV concurrency drops a hair (7.65× → 7.33× @ 256K; capture reserves a bit more workspace) — a trivial cost.

(Expert parallelism, --enable-expert-parallel, was tested and rejected: the W4A16 experts are already 4-bit-tiny, so sharding them frees ~nothing — same concurrency, same/worse tok/s, only faster load.)

The result

  • 69GB W4A16, group_size=32, compressed-tensors / Marlin (down from 219GB BF16).
  • Serves on 8× RTX 3090, TP=8, --gpu-memory-utilization 0.90, --max-model-len 262144, no --enforce-eager, poolside_v1 tool/reasoning parsers, enable_thinking.
  • ~92 tok/s single-stream decode, 7.33× concurrency at full 256K.
  • Clean output with thinking on, including technical/markdown content.
  • Terminal-Bench (terminal-bench-core 0.1.1, 18-task subset, terminus agent, pass@1): 4/18 = 22.2% — solved crack-7z-hash.easy, sqlite-with-gcov, prove-plus-comm, sanitize-git-repo; near-solve on swe-bench-astropy-2 (8/9 sub-tests). The hard infrastructure tasks (QEMU, kernel-build, maze, video, chess, password-recovery) hit the 900s agent timeout — expected for an 8B-active model on a frontier-calibrated benchmark.

How to run

Quantize (box with the BF16 base + enough RAM for one MoE layer; disk offload handles the rest):

~/awq-quant/bin/python w4a16-quant-laguna-s-2.1-perlayer.py

Serve (8× 24GB Ampere):

bash Laguna-S-2.1-W4A16-vLLM.sh

vLLM 0.25.1+, --quantization compressed-tensors (auto-detected from the model's quantization_config), poolside_v1 parsers, enable_thinking, TP=8, 0.90 util, 256K, NO --enforce-eager (cudraphs is the 12× lever). Full flags in the script.

Files

file what
w4a16-quant-laguna-s-2.1-perlayer.py the per-layer linearize→quantize loop (Wall 2) with group_size=32 (Wall 4)
Laguna-S-2.1-W4A16-vLLM.sh the vLLM serve recipe (cudagraph, Wall 5)
LAGUNA-W4A16-BLOG.md this writeup

Why this matters

Laguna S 2.1 is a top open agentic-coding MoE that fits consumer hardware — but only if you can quantize it without OOMing, get the group_size right, and serve it without leaving throughput on the floor. The per-layer linearize trick generalizes to any 256-expert-class MoE that OOMs llmcompressor's global linearize. The group_size lesson — match the reference quant's granularity, don't inherit it from a different architecture — generalizes to every MoE quant. And on Ampere, measure with cudagraphs on before you conclude a model is slow.


Built on an 8×3090 rig (192GB VRAM, ~251GB RAM). Quant env: llmcompressor 0.12 on transformers 5.10.1. vLLM 0.25.1. Model: OpenMDW-1.1 (poolside). Quant scripts are MIT-adjacent — use them freely.