jarrelscy's picture
Publish GLM-5.3 NVFP4 / ARVQ v2 same_input fp16-block hybrid: non-cold structure seed
7315531 verified
|
Raw History Blame
17.7 kB

ARVQ cold-expert quantization: methodology

How the cold Mixture-of-Experts weights in this repository (GLM-5.3-Vision-NVFP4-ARVQ-hybrid) were quantized and tuned. This document reconstructs the actual pipeline and tuning decisions from the build transcripts and the fitting/serving code; numbers are cited from those sources. It is a methodology record, not a quality claim: full-model task quality and native SM120 execution were not evaluated for this checkpoint.


1. What this model is (the hybrid)

GLM-5.3 is a large MoE. Each MoE block routes each token to its top-8 of 256 routed experts (plus a shared expert). This checkpoint is a hybrid in which each expert is stored in one of two ways:

  • Hot experts — kept in the NVFP4 donor format (unchanged).
  • Cold experts — re-quantized to ~2 bits with ARVQ (additive residual vector quantization), the subject of this document.

The hot/cold split is fixed by a REAP allocation that keeps the 5,750 most important experts hot; the remaining 13,450 experts across the model are cold (per the 09-14 build transcript, which verified "13,450 cold-expert source files"). The allocation is 75% text-REAP / 25% multimodal-salience weighted and is not modified by this campaign.

Only the 75 MoE blocks, layers 3–77, are touched. Layers 0–2 are dense and stay frozen. Everything except the cold experts is inherited unchanged: attention, backbone, shared experts, hot NVFP4 experts, BF16 MTP, and the vision components.

Cold experts are quantized per (layer, projection). The two projections are the fused gate/up w13 (N=4096, K=6144) and the down projection w2 (N=6144, K=2048).


2. ARVQ representation and bit budget

Each cold expert weight group of 8 contiguous columns is represented as the sum of two codebook atoms times a scale:

w_group = global * block_scale * (c0[a] + c1[b])
  • Two codebooks c0, c1, each 256 entries × 8 dims (512 codewords per pair). This is the "rvq256_256x8" format. Codebooks are per-expert in this v3 checkpoint (codebook_scope = expert, packed format rvq256_256x8_expert, version 3).
  • Indices a, b are one uint8 each per 8-weight group → 16 index bits per group = 2.0 bits/weight.
  • Codebook atoms are constrained to the FP4 grid {0, ±0.5, ±1, ±1.5, ±2, ±3, ±4, ±6} (project-to-FP4 after every update). The FP4 grid on the atoms is why no incoherence/Hadamard rotation is used: the target space cannot be rescaled without leaving the grid (see encoder.py docstring).
  • Block scales are stored as float8_e4m3fn, one scale per [row, 128-column block] (48 blocks/row for w13, 16 blocks/row for w2). At use time each block scale is repeat_interleave(16) to cover the 16 eight-wide groups in a 128-column block. Adding the FP8 scale (8 bits / 128 weights = 0.0625 bit/weight) gives the reported 2.0625 bits/weight plus the amortized per-expert codebooks.
  • A single fp32 global scalar per (layer, projection) sits above the block scales so the index/scale layout is unchanged from earlier versions.

Packing (pack.py) writes MMA-fragment index tiles (64 uint32 words/tile) and verifies, per layer, that every expert's indices round-trip exactly and that a sample of decoded weights matches bit-for-bit before publishing.


3. Initial fit (Hessian-aware, no PV yet)

Before any gradient tuning, each (layer, projection) gets an activation-Hessian-tuned initialization (encoder.py, fit.py):

  1. Per-expert raw-basis Hessian. w13 uses H = E[x xᵀ] over routed tokens; w2 uses H = E[m mᵀ] where m = silu(x·Wgᵀ)·(x·Wuᵀ) is the BF16 SwiGLU intermediate. Experts with <32 routed tokens fall back to all captured tokens. Activations are stored/computed in float32; Hessians are half.
  2. Escalating-damp Cholesky of H⁻¹: damping starts at 1e-3 · mean(diag H) and multiplies by 4× (up to 12 attempts) until the factorization succeeds.
  3. Codebook EM on a Hessian-importance-weighted subsample (20,000 groups/ expert) of normalized 8-dim groups: k-means warm start → 8 alternating refine iterations, FP4-projecting the atoms after every M-step.
  4. Per-expert LDLQ/GPTQ error-feedback column sweep with the fixed codebooks (col_block=128, 2 sweep passes, 1 inner refine): plain-L2 inner assignment, off-diagonal Hessian carries cross-group error feedback.
  5. Least-squares refit of the E4M3 block scales given the chosen codes.

An earlier decision, recorded 09-14: a tuned Hadamard rotation was tested and rejected — a matched 48-expert down-projection refit gave 0.372 (H128 rotated) versus 0.191 unrotated, so the unrotated raw-basis fit was adopted.

Nonfinite guards are raised as errors throughout (the fit/tuner refuse to proceed on NaN/Inf rather than silently scrubbing).


4. The PV objective: --target reference

The cold experts are then output-tuned per layer. The objective for this checkpoint is --target reference (sequential_pv_full_corpus.py).

For each captured token position the tuner forms:

  • frozen — the block output from everything that is not a cold expert on the student input (retained tuned upstream layers → attention → shared expert → hot NVFP4 experts). This is fixed.
  • reference — the original FP8 reference block output: the unchanged donor backbone with the original FP8 routed experts (all 256), evaluated on the original reference trajectory for that token position.
  • required = reference − frozen — the residual the cold experts must supply.

The loss is the relative squared error of the cold-expert sum against required:

loss = || student_cold(x_student) − required ||²  /  (denominator · Σw)

where denominator is the per-row mean energy of reference − frozen over the corpus (so the scale is comparable across layers). Selection metric is reference_rel = ||(frozen + cold) − reference|| / ||reference||, measured over the whole transformer block including the residual, not the cold-experts sum alone.

Why reference (and its known weakness). The transcripts weighed reference against a same_input target (match the block's output on the student's own drifted input). The rationale for reference (09-16): "Matching the original reference trajectory is more directly aligned with preserving the model's behavior." The acknowledged risk: because each layer receives a different (student) input than the original, the cold experts may lack the capacity to reproduce the reference output and could "learn corrections that work on training examples but fail on unseen ones" — i.e. reference-target is the more out-of-distribution objective. It was chosen only after a matched layer-4 A/B:

Layer-4 metric reference vs same_input
validation reference error 0.21% better with reference
development-audit error 0.08% worse with reference

A small mixed difference; reference was retained for behavioral alignment. Two important design consequences: the target is the complete block output (including the residual), so matching only the MoE contribution cannot leave the incoming residual error untouched; and the loss compares the routing-weighted sum of all cold experts jointly (gate/up and down together).


5. Sequential per-layer pipeline with trajectory coupling

Layers are processed sequentially, 3 → 77 (full_reference.py / sequential_capture.py):

  1. Capture (8-way sequence-parallel). For layer L, the student input is the retained tuned output of block L-1 (student_outputs.pt), while the reference target re-runs the donor backbone + original FP8 experts on the original reference hidden state. Both trajectories are carried forward in a rolling cache; layer 3 is the control (no preceding ARVQ layer, so student = reference input). This cross-layer coupling means each layer is tuned against the exact drifted inputs it will see at serving time.
  2. Fit / PV tune the cold experts (Section 6).
  3. Gate (Section 7).
  4. Export + serialized replay — decode to the packed format and require the reloaded forward to match the tuned forward to <2e-6 before use.
  5. Per-layer HF publish — publish_reference.py re-assembles cold slots across the 8 expert shards, checks slot coverage/allocation/global metadata, runs the exact index round-trip and sampled decoded-weight parity, then atomically replaces that layer's two tensor files + reports in one commit with parent-commit protection; remote hashes are re-verified.

Parallelism: cold experts are sharded across 8 GPUs (each rank owns distinct experts); routed outputs are summed with an all-reduce so a single coupled-output objective is optimized without a dense-model replica.


6. PV tuning numerics (the published full-corpus campaign)

Per the model card, all 75 layers were replaced by a full-corpus sequential PV campaign drawing sequentially from 18,001,846 training tokens at context 1024 (53,248 legacy tokens were removed around 17 exact matches to held-out prompts). Fixed validation and development-audit sets each hold 16,384 tokens.

  • Optimizer: Adam. Two trainable parameter groups — FP4-constrained per-expert codebooks and FP8-constrained per-block scales. (The faster fork trains per-row scale deltas instead; the published full-corpus fork trains the per-128-block log-scales.)
  • Effective batch: 262,144 tokens, accumulated as four 65,536-token microbatch passes (256 whole 1024-token sequences), one shuffled no-replacement pass over the corpus.
  • Learning-rate schedule: book/scale LR 0.048 / 0.032 through update 45, then dropped by ×0.25 to 0.012 / 0.008 (lr_decay_after=45, lr_decay_factor=0.25). (Note: these campaign LRs are higher than the in-repo code defaults of 0.003/0.002, consistent with the much larger 262k-token batch.)
  • Update budget: layers 4–26 use a fixed 69 updates. From layer 27, 69 is the maximum with an adaptive early stop: three validation checks more than 0.1% worse than best trigger an earlier LR reduction; 15 updates without improvement after that reduction permit stopping.
  • Validation every 5 updates, retaining the best checkpoint.
  • Index reassignment every 20 updates (Section 6.1).
  • Regularization: codebook drift-from-anchor penalty + scale drift penalty, weighted 0.01; grad-norm clip 1.0; codebook atoms clamped to [-6, 6]; log-scales clamped to within ±0.35 of their initial value.
  • Coverage gate: experts with fewer than 256 routed training rows are frozen (gradient masked) and keep their initial books/scales.

6.1 Output-gradient index reassignment

Continuous Adam cannot move the discrete indices, so every 20 updates a discrete proposal pass runs (gradient_indices.py, method expert_parallel_output_gradient_prefix_backtracking_v1):

  1. Backprop the output loss to the reconstructed weight of one expert projection.
  2. Propose alternative (a,b) index pairs near a small gradient step, capped at max_fraction=0.001 of groups, trust_ratio=0.01, target_ratio=0.03; the current pair is always in the candidate set (changing nothing stays legal).
  3. Prefix backtracking acceptance: try the top-k proposals with k∈{full, ¼, 1/16, 1}, accept only if the training SSE strictly drops and a separate check batch SSE does not increase (beyond 1e-7 slack); otherwise revert. Proposal deltas are broadcast sparsely (only routed rows).

The transcripts flag this as the method's main theoretical weakness versus published PV-Tuning: codebooks are tuned for output accuracy while index proposals originally came from weight-reconstruction Hessians — the output-gradient proposal above was added to close that gap.

6.2 Cold arithmetic emulation

Both training and evaluation emulate the serving numerics (activation.py, pinned to serving revision b1380cf7…): four FP4 activation planes (fp4_planes4_fp16_boundaries_v1) and FP16 SwiGLU boundaries (FP32 GEMM, FP16 SiLU/product), with straight-through estimators for gradients. This emulates the quantization boundaries only — native SM120 MMA accumulation and TP reduction are not bit-exact and were not qualified.


7. Acceptance gates

Publication of a layer requires all of:

  • Validation non-regression: 0 ≤ final_reference_rel ≤ initial·(1+1e-6).
  • Development-audit non-regression: the held-out audit split (never used for updates or checkpoint selection) must satisfy audit_reference_rel ≤ initial_audit·(1+1e-6); otherwise the layer keeps its initialization. There is no absolute error floor — the rule is purely final ≤ initial.
  • Serialized-replay parity: decoded/reloaded forward matches the tuned forward to <2e-6, and remote file hashes are re-verified after upload.

The audit is honestly labeled development data, not an untouched final test; it is a split of the same calibration tokens the initialization already saw. Two notable per-layer decisions recorded in the card: layer 30 retains its audit-qualified update-15 checkpoint after its validation-best update-20 failed the development audit; layer 3 retains a separately qualified lower-LR refinement.


8. Calibration corpus

The text corpus (build_calib_v31.py + reasoning_slice.py) is deterministic (seed 42), GLM-tokenized, with the following domain mix by tokens:

Share Domain Sources
~32% code local vLLM sources, m-a-p/CodeFeedback, jtatman/python-code-500k
~20% tool-calling / agentic generated GLM chat-template sessions
~15% reasoning OpenR1-Math-220k + dolphin-r1 <think>…</think> → answer
~13% instruction chat tatsu-lab/alpaca + coding chat
~10% medical MedQA textbook continuation + medical Q&A
~10% prose vLLM docs markdown + databricks-dolly-15k

The reasoning slice was added specifically because the earlier calib-v3 mix starved the cold 2-bit experts of reasoning-termination behavior: 87% of its <think> blocks were empty. The slice emits multi-turn conversations where an earlier assistant turn carries a real <think>…</think> that terminates and hands off to an answer, followed by a trailing user turn, so the think-close token </think> (id 154842) renders in-stream rather than as a trailing generation prompt. Held-out prompts are explicitly excluded (last 25 vLLM docs, last 3 MedQA files, vLLM code beyond index 400 of the seed-42 shuffle), and 17 exact-match prompts were purged from the training set.

The capture/training code also supports up-weighting boundary rows (rows whose next token is a boundary such as </think> / <|endoftext|>) via a row_weight term whose sum normalizes the loss, and an optional matched-mixed validation set; the published card describes reasoning-based validation.


9. Results captured during the build

All errors below are held-out relative L2, either over the whole transformer block (including residual) or over the cold-experts sum, as noted. They are local reconstruction errors, not token-accuracy or perplexity.

  • Pilot smoke test (layer 3, one step): full-block error 0.005497 → 0.005456 (~0.75%), 304/388 expert-projection proposals accepted; exported weights reproduced the retained result.
  • Layer 3 (200-step pilot): full-block reference error 0.005497 → 0.005271 (−4.12%), best at step 200.
  • Layer 4 (200-step pilot): full-block reference error 0.015760 → 0.015596 (−1.05%).
  • reference vs same_input (layer 4, matched): reference 0.21% better on validation, 0.08% worse on development audit.
  • Faster fork parity: the optimized fork (larger microbatch, specialized embedding-lookup backward, sparse proposal broadcast) matched the reference fork's validation and audit exactly on layers 3 and 4, with all propagated BF16 outputs equal; training loops fell from ~1525 s → ~337 s (layer 3) and ~1226 s → ~303 s (layer 4).
  • Early-layer stability example (layer 18, full-corpus): initial 0.017205, step-5 0.017176, step-69 0.017183 — differences ≤0.17%, treated as a signal to watch LR/noise rather than proof of convergence.

Older cold-only checkpoints reported ~0.2211→0.2070 (layer 3) and 0.2364→0.2320, but those measured cold-expert output only on different captures and are not comparable to the full-block errors above.

No comparison against the FP8 donor or an AQLM variant on perplexity / KLD / top-1 / wikitext was recorded in the reviewed transcripts, and no full-model task evaluation was performed. Those remain open.


10. Honest limitations

  • Full-model quality and native SM120 execution were not evaluated; the arithmetic emulation covers FP4/FP16 boundaries only, not MMA/TP bit-exactness.
  • The development audit is a split of calibration data, not an independent test.
  • The reference target is the more OOD objective; its generalization advantage over same_input was small and mixed on the one matched layer tested.
  • Gates enforce local non-regression, not any absolute quality bar.

Prepared from the GLM-5.3 ARVQ build transcripts (2026-09-14 and 2026-09-16) and the btx53/arvq88 fitting/serving code. Where a number could not be sourced it is stated as unknown rather than estimated.