Upload folder using huggingface_hub

fdc6474 verified 3 months ago

4.02 kB

# `aqlm_converge.py` — activation-aware AQLM convergence for cold experts
## What it does
The shipped checkpoints quantized cold experts with **init-grade** AQLM:
residual k-means in *weight space* (minimize ‖W−Ŵ‖²), which ignores that
some weight columns matter far more than others at inference time. This
tool re-optimizes those same codes/codebooks/scales against the objective
that actually matters — **per-expert output error on real routed traffic**:
```
minimize  Σ_e ‖(W_e − Ŵ_e) · diag(√h_e)‖²_F
```
where `h_e = E[x²]` is a per-expert **diagonal input Hessian** estimated
from calibration activations that were actually routed to expert `e`:
- for **w13** (gate/up): `x` = the layer's hidden-state inputs
- for **w2** (down): `x` = that expert's own `silu(gate(x))·up(x)`,
  computed with **teacher** (original NVFP4-dequantized) weights
Intuition: input dimensions that carry large activations get their weights
approximated more carefully; near-dead dimensions absorb the quantization
error. This is the same idea as GPTQ/AQLM's calibration objective, with a
diagonal (rather than full) Hessian so every step stays a giant batched
GEMM that a B200 executes in seconds.
### Algorithm (per layer × projection, over the layer's cold experts)
1. **Warm start** from the shipped codes/codebook/scales (already a decent
   weight-space solution — convergence needs few iterations).
2. Iterate (max 6, early-stop when improvement < 0.2%):
   a. **Weighted coordinate-descent re-encode** — for every group of 8
      weights pick the codebook entry minimizing the h-weighted error.
      Cost per sweep = two GEMMs against the 65,536×8 codebook.
   b. **Closed-form weighted codebook update** — each entry ← the
      h-weighted mean of its assigned groups (weighted k-means step);
      dead entries respawn on random groups.
   c. **Closed-form weighted scale refit** — per-expert per-out-channel
      least squares under h.
3. Log the weighted relative error before/after (typical: w13 ~0.09 →
   substantially lower on hot input dims; the *unweighted* MSE may barely
   move — that's the point).
What this still is NOT: full AQLM beam search (beam=1 CD only) and no
PV-tuning/global finetune. It captures the largest quality lever
(activation-aware objective) at ~1/100th the cost.
## Prerequisites (all already produced on this box)
| input | path | produced by |
|---|---|---|
| calibration activations, ~24k routed tokens/layer | `/data/glm52-acts/acts_layer{N}.pt` | `capture_acts.py` (needs the `VLLM_ACT_CAPTURE_DIR` hook, vLLM patch ≥0002) |
| teacher weights (original per-expert NVFP4) | `/tmp/glm52-hot-dl2` regions + `/data/glm52-old-layerwise` (layers 3,4,5,8,74–77) | `make_hot_manifest2.py` + `range_download.py` |
| warm-start codes + cold-expert ids | `/data/glm52` (live two-tier checkpoint) | — |
## Running
```bash
cd /home/coder/git/glm52 && source .venv/bin/activate

# smoke test (2 layers, idle GPUs):
python tools/aqlm_converge.py --layers 40,76 --gpus 4,5
# full run (75 layers, all 8 GPUs, layer-parallel):
python tools/aqlm_converge.py > /data/glm52-aqlm-conv/run.log 2>&1 &
```
Resumable: finished layers (`/data/glm52-aqlm-conv/layer_N.pt`) are
skipped on rerun. Expect roughly 10–25 min/layer/GPU (dominated by the
fp32 error/dequant temporaries — each worker peaks ~170 GB GPU memory, so
one layer per GPU at a time).
## Output & downstream
`/data/glm52-aqlm-conv/layer_{N}.pt`:
`expert_ids [nC]`, `w13_codes [nC,1,4096,768] i16`, `w13_codebooks
[1,65536,8] f16`, `w13_scales [nC,4096] f16`, `w2c_*` likewise, plus
`*_err_before/after` (h-weighted rel-err).
To ship: rebuild the checkpoints' cold arrays from these parts (slice by
`expert_ids` per target's cold set — the 1M cold set is a superset of the
500k/250k ones), revalidate (coherence + needle + KV/memory), re-upload.
The checkpoint format, kernels, and serving recipe are unchanged — only
the bytes inside `w13_*` / `w2c_*` tensors improve.

	# `aqlm_converge.py` — activation-aware AQLM convergence for cold experts

	## What it does

	The shipped checkpoints quantized cold experts with init-grade AQLM:
	residual k-means in weight space (minimize ‖W−Ŵ‖²), which ignores that
	some weight columns matter far more than others at inference time. This
	tool re-optimizes those same codes/codebooks/scales against the objective
	that actually matters — per-expert output error on real routed traffic:

	```
	minimize Σ_e ‖(W_e − Ŵ_e) · diag(√h_e)‖²_F
	```

	where `h_e = E[x²]` is a per-expert diagonal input Hessian estimated
	from calibration activations that were actually routed to expert `e`:

	- for w13 (gate/up): `x` = the layer's hidden-state inputs
	- for w2 (down): `x` = that expert's own `silu(gate(x))·up(x)`,
	computed with teacher (original NVFP4-dequantized) weights

	Intuition: input dimensions that carry large activations get their weights
	approximated more carefully; near-dead dimensions absorb the quantization
	error. This is the same idea as GPTQ/AQLM's calibration objective, with a
	diagonal (rather than full) Hessian so every step stays a giant batched
	GEMM that a B200 executes in seconds.

	### Algorithm (per layer × projection, over the layer's cold experts)

	1. Warm start from the shipped codes/codebook/scales (already a decent
	weight-space solution — convergence needs few iterations).
	2. Iterate (max 6, early-stop when improvement < 0.2%):
	a. Weighted coordinate-descent re-encode — for every group of 8
	weights pick the codebook entry minimizing the h-weighted error.
	Cost per sweep = two GEMMs against the 65,536×8 codebook.
	b. Closed-form weighted codebook update — each entry ← the
	h-weighted mean of its assigned groups (weighted k-means step);
	dead entries respawn on random groups.
	c. Closed-form weighted scale refit — per-expert per-out-channel
	least squares under h.
	3. Log the weighted relative error before/after (typical: w13 ~0.09 →
	substantially lower on hot input dims; the unweighted MSE may barely
	move — that's the point).

	What this still is NOT: full AQLM beam search (beam=1 CD only) and no
	PV-tuning/global finetune. It captures the largest quality lever
	(activation-aware objective) at ~1/100th the cost.

	## Prerequisites (all already produced on this box)

	\| input \| path \| produced by \|
	\|---\|---\|---\|
	\| calibration activations, ~24k routed tokens/layer \| `/data/glm52-acts/acts_layer{N}.pt` \| `capture_acts.py` (needs the `VLLM_ACT_CAPTURE_DIR` hook, vLLM patch ≥0002) \|
	\| teacher weights (original per-expert NVFP4) \| `/tmp/glm52-hot-dl2` regions + `/data/glm52-old-layerwise` (layers 3,4,5,8,74–77) \| `make_hot_manifest2.py` + `range_download.py` \|
	\| warm-start codes + cold-expert ids \| `/data/glm52` (live two-tier checkpoint) \| — \|

	## Running

	```bash
	cd /home/coder/git/glm52 && source .venv/bin/activate

	# smoke test (2 layers, idle GPUs):
	python tools/aqlm_converge.py --layers 40,76 --gpus 4,5

	# full run (75 layers, all 8 GPUs, layer-parallel):
	python tools/aqlm_converge.py > /data/glm52-aqlm-conv/run.log 2>&1 &
	```

	Resumable: finished layers (`/data/glm52-aqlm-conv/layer_N.pt`) are
	skipped on rerun. Expect roughly 10–25 min/layer/GPU (dominated by the
	fp32 error/dequant temporaries — each worker peaks ~170 GB GPU memory, so
	one layer per GPU at a time).

	## Output & downstream

	`/data/glm52-aqlm-conv/layer_{N}.pt`:
	`expert_ids [nC]`, `w13_codes [nC,1,4096,768] i16`, `w13_codebooks
	[1,65536,8] f16`, `w13_scales [nC,4096] f16`, `w2c_*` likewise, plus
	`*_err_before/after` (h-weighted rel-err).

	To ship: rebuild the checkpoints' cold arrays from these parts (slice by
	`expert_ids` per target's cold set — the 1M cold set is a superset of the
	500k/250k ones), revalidate (coherence + needle + KV/memory), re-upload.
	The checkpoint format, kernels, and serving recipe are unchanged — only
	the bytes inside `w13_` / `w2c_` tensors improve.