METHODOLOGY.md

# 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.*