docs/FINDINGS.md

# Step-3.7-Flash REAP: Measured Findings

**Model:** Step-3.7-Flash — StepFun 198B VLM (196B language backbone + 1.8B ViT),
45 decoder layers (dense 0–2, **MoE 3–44 = 42 layers**), **288 experts/layer**,
top-k=8, 1 shared expert, sigmoid router with `e_score_correction_bias`, SwiGLU,
RMSNorm. 375 GB BF16 on disk. Experts = **97% of parameters**.

**Hardware:** M5 Max, 128 GiB unified memory (`iogpu.wired_limit_mb` = 115 GiB).

**Scope:** Everything below is *measured*, not asserted. Where a conclusion was
later overturned by better data, both the wrong and corrected versions are shown —
several "obvious" fixes measured *worse* than what they replaced.

---

## TL;DR — the three measured effects, ranked

| # | Finding | Magnitude | Status |
|---|---|---|---|
| 1 | **Vision blindness** — text-only saliency can't see vision experts | **9.42%** of vision saliency mass sits on pruned experts; verdict **SEPARATE** (ρ=0.381) | ⚠️ **unaddressed** |
| 2 | **Layer-adaptive allocation** — uniform per-layer keep counts are suboptimal | **5.3%** less discarded mass at p15, free from existing data | available, unbuilt |
| 3 | **Truncation fix** — 384-token head truncation dropped answers | **1.7%** of experts changed in the final plan | ✅ done, minor |

**The vision issue is the largest open problem by a wide margin.**

---

## 1. The reap decision: 15% (keep 245/288)

Discarded router-saliency mass by rung (5k prompts; 1k and 5k agreed to within
0.05%, so **saliency stabilises by ~1k prompts**):

| Reap | Keep/288 | Discarded mass (mean / worst layer) | 4-bit size |
|---|---|---|---|
| 10% | 260 | 3.0% / 4.0% | ~89 GB |
| **15%** | **245** | **5.0% / 6.8%** | **~84 GB** |
| 20% | 231 | 7.1% / 9.6% | ~80 GB |
| 25% | 216 | 9.6% / 12.9% | ~75 GB |

**This model resists pruning.** Every expert fires (zero dead experts), saliency is
flat, and prune sets are layer-local (Jaccard ~0.12 across depth — only ~4 experts
prunable everywhere). Unlike models where 25–50% comes free, capacity here is real.

Rationale for 15%: it's the only rung with empirical validation (a prior 158B ≈ 15%
build worked well); cost accelerates past it; and **reaping is not the size lever** —
quantization is. 15%→20% saves ~4 GB while nearly doubling worst-layer loss.

Raw REAP scores grow ~140× with depth (mean 0.17 at layer 3 → 24.0 at layer 44), so
**scores are not comparable across layers** without normalization.

---

## 2. Vision blindness ⚠️ — the biggest open issue

### The mechanism (verified in source)
Image features are `_masked_scatter`'d into `inputs_embeds` (`step3p7.py:105-111`)
and flow through the **same 42 MoE layers** being pruned. There is no separate
"vision MoE" — image tokens are routed by the same routers to the same experts.

But the collector builds embeddings via `text.embed_tokens(tokens)` only
(`collect_step3p7.py:312`) — it never calls `vision_model` or
`get_input_embeddings`. **So experts specialising in image tokens receive zero
activation during saliency collection, score ~0, and are pruned first.**

Net effect: the pipeline carefully preserves the vision *encoder* at BF16 while
deleting the experts that *consume its output*.

### The experiment
A dedicated vision-only saliency pass (`reap_stream/collect_vision_saliency.py`)
routed **300 images through the real vision path** — processor → `pixel_values` →
`get_input_embeddings` → merged embeds → the same block-streaming loop. Calibration
was deliberately **mixed: 200 ChartQA (synthetic plots) + 100 VQAv2 (natural
photos)**, because charts alone are a narrow slice of "vision" and would risk a
falsely reassuring result. Runtime: 1302 s.

Verdict thresholds were **pre-committed** before seeing data (MIXED if ρ>0.75 and
overlap>0.6; SEPARATE if ρ<0.4 or overlap<0.35).

### Result: SEPARATE

| Metric | Value |
|---|---|
| Mean rank correlation (vision vs text) | **0.381** |
| Mean top-25% expert overlap | **0.478** |
| Vision-top experts inside p15 prune set | **3.1 per layer** (max 9) |
| **Vision saliency mass on pruned experts** | **9.42%** (max **13.40%**) |

Correlation is lowest in early layers (~0.16–0.33, layers 3–14) and rises to
~0.36–0.54 through mid and late layers. The effect is **roughly uniform across
depth** — it is not concentrated anywhere in particular.

> **Correction — an error in earlier analysis.** An initial run of this comparison
> used the *partial* checkpoint, where layers 36–44 had `reap_count = 0`
> (not yet collected). Correlating against all-zero text saliency produced
> meaningless near-zero/negative values (e.g. layer 37: −0.145, layer 40: −0.115),
> which were mistakenly reported as a real "late-layer collapse" and called the most
> interesting finding. **On complete data those layers are normal** (layer 37: 0.457,
> layer 40: 0.435, layer 44: 0.537). There is no late-layer collapse. The headline
> verdict (SEPARATE) and the 9.4% mass figure were unaffected — both were robust
> across partial and complete data.

### Options (neither built)
1. **Protect vision-heavy experts** — union the top vision-saliency experts per layer
   into the keep set before applying the ratio cut. Cheap; slightly raises expert count.
2. **Combine explicitly via Contrastive Skill-Shield** (Theorem 6 in
   `frontier-moe-pruning-theorems.md`), treating vision as the target domain:
   `S_CS(e) = S̃_text(e) + α·max(0, S̃_vision(e) − S̃_text(e))`.
   **z-score normalization is required** — vision came from 300 images vs text's 2500
   prompts, so raw magnitudes aren't comparable.

Keeping the two saliency maps **separate and combining explicitly** is preferable to
mixing vision into the calibration set: mixing makes the blend ratio an implicit
hyperparameter and destroys the ability to see which modality wanted which expert.

---

## 3. Calibration data defects (both real, both fixed)

### 3a. Truncation dropped the answer
The collector did `tokenizer.encode(text)[:384]`. Share of prompts whose
`ASSISTANT:` response starts **beyond** the token budget:

| Category | beyond 384 | beyond 1024 |
|---|---|---|
| reasoning_math | 0.4% | ~0% |
| coding | 8.8% | 2% |
| tool_use | 41% | 1% |
| general_instruction | **96.3%** | 35% |
| **agentic** | **96.4%** | **64%** |

Token distribution: **median 798, p90 5195, p99 6380, max 10,182.** Agentic prompts
(system prompt + tool schemas + file contents) are enormous — no front-truncation
budget reaches their answers.

**Fidelity measured against full-length ground truth** (not assumed):

| Window / mode | Jaccard vs full | Spearman vs full |
|---|---|---|
| 384 head (original) | 0.783 | 0.938 |
| 1024 head | 0.834 | 0.960 |
| 1024 **tail** | **0.694** ← worst | 0.917 |
| **1024 headtail** | **0.858** ← best | **0.977** |

> **`tail` — the initially proposed fix — measured WORST.** Most routing statistics
> come from processing the long context, not the short answer; tail truncation
> discards that context. Tested before shipping; the intuition was wrong.

`_truncate()` now supports `head|tail|headtail` via `--truncation`.

### 3b. Multimodal rows were misaligned
`build_calib_mix.py:186` paired local PNGs (sorted by filename) against a **freshly
re-loaded** HF ChartQA dataset **by raw index**, silently desyncing image and Q&A.
Verified on two samples (an antiretroviral-therapy chart paired with a "positive
view peak" question; a Trump-approval chart paired with an online-classes question).
All 500 multimodal rows suspect.

- Excluded from DWQ via `artifacts/dwq-targets/exclude_indices.json` (105 of 1500).
- Fixed properly by `scripts/build_multimodal_calib_fixed.py` /
  `build_multimodal_calib_mixed.py`, which read image + query + answer from the
  **same dataset row in one pass** so they cannot desync.

**Note:** the misalignment was irrelevant to *REAP saliency* (no images were ever
fed), but genuinely harmful to *DWQ*, where the text-only teacher was asked to
predict chart answers it could not see.

---

## 4. Result of the corrected saliency re-run

Re-ran with `--max-tokens 1024 --truncation headtail --max-samples 2500`
(1.40 h, 42 layers, `nested_ok: true`) → `artifacts/step37-1024-headtail/`.

**Plan diff, new p15 vs original p15:**

| Metric | Value |
|---|---|
| Experts changed per layer | **4.3 of 245 (1.7%)** |
| Range | 2–10 |
| Total slots changed | 180 across 42 layers |
| Prune-set Jaccard (old vs new) | **0.821** |

**~98% of the pruning decision was identical.** The truncation flaw was real and
correctly diagnosed, but its downstream effect on the plan is small — expert ranking
is robust to this input-distribution shift. Changes cluster in early layers
(layer 3: 9 changed) and are smallest late (layers 36–44: ~3), consistent with early
layers being more sensitive to surface/formatting features.

Use the new plan going forward (it is more faithful and already paid for), but it is
**not** a meaningful quality upgrade on its own.

---

## 5. Layer-adaptive expert allocation (measured, unbuilt)

Current plans keep a **uniform 245 experts in every layer**. Allocating the same
total budget non-uniformly — greedily, by *fractional* mass cost per layer since raw
scores aren't cross-comparable — measurably reduces total discarded saliency:

| Rung | Uniform loss | Adaptive loss | Improvement | Adaptive keep range |
|---|---|---|---|---|
| p15 | 2.1107 | 1.9996 | **5.3%** | 190–276 (mean 245) |
| p25 | 4.0473 | 3.9442 | 2.5% | 150–262 (mean 216) |

Computed from **existing saliency** — no new collection needed. Caveats: discarded
mass is a proxy, not a quality measurement; and adaptive allocation **breaks the
nesting property** (p10 ⊂ p15 ⊂ p20 ⊂ p25) that the launcher currently verifies.

---

## 6. Memory: the biggest operational lesson

### MLX under-reports GPU memory by >2×
The collector logged `peak_mb = 46506` (~46 GB) while the process actually held
**109–110 GB**:

```
footprint <pid>
  109 GB   IOAccelerator (graphics)   [all DIRTY, 0 reclaimable]
  phys_footprint: 110 GB
```

**Always use `footprint <pid>`, never `mx.get_peak_memory()`.** Trusting MLX's number
allowed two crashes to happen unseen. (Note: `pgrep -f run_mac_bf16` also matches the
zsh wrapper — filter for the real Python PID.)

### Failure modes observed
1. **OS-level hard reboot** — the whole machine went down mid-run.
2. **GPU watchdog timeout** — `[METAL] Command buffer execution failed: Caused GPU
   Timeout Error`. Memory pressure stalls GPU ops until macOS kills the buffer.

**Leading indicator for both: `vm_stat` compressor > 40 GB and swap expanding.**
Healthy = compressor ≈ 0. Do *not* watch "free RAM" — macOS always reports ~0.

### Root cause and fix
Not `layers_at_once`, not prompt count. **MLX's allocator hoards freed buffers**, and
`mx.clear_cache()` was only called once per *layer* — so 2500 per-prompt forwards
accumulated every freed intermediate.

```python
for i in range(len(hidden)):
    hidden[i] = _run_layer(layer, hidden[i], sliding_window)
    if (i + 1) % _CACHE_EVERY == 0:      # _CACHE_EVERY = 200
        mx.clear_cache()
```

| | Before | After |
|---|---|---|
| IOAccelerator | **109 GB** | **34 GB** |
| Compressor | 70 GB | **0.1 GB** |
| Swap used | 9.6 GB | 553 MB |

**~3× reduction at zero cost** to sample count, layers, or plan quality. Also fixed:
`hidden = [_run_layer(...) for h in hidden]` materialised an entire second ~21 GB list
before releasing the old one — now updates in place.

With this fix, the text (≈48 GB) and vision (≈14 GB) passes ran **concurrently** with
compressor at only 1.0 GB.

### Sizing reference
- MoE layer weights: **9.06 GB each** (288 × 3 proj × 4096×1280 × BF16)
- Hidden states: `n_prompts × max_tokens × 4096 × 2 bytes` (2500×1024 ⇒ ~21 GB)
- `iogpu.wired_limit_mb` is already 115 GiB; **raising it further starves the OS** and
  causes the hard-reboot failure mode.

---

## 7. DWQ: built, works, value unproven

Pipeline (`mlx_vlm`-native, since `mlx_lm` only knows `step3p5`, not `step3p7`):
- **Phase 1** `dwq_collect_targets.py` — streams the BF16 teacher, caches **top-128**
  logits per prompt (377 MB for 1500 prompts, ~32 min). Full-vocab would be ~495 GB
  and is unnecessary: LM distributions are peaked, top-128 captures >99.9% of mass.
- **Phase 2** `dwq_train_student.py` — 92 GB student resident, trains only affine
  quant scales, checkpoints the small (~5.3 GB) trainable tensors with auto-resume.

### Blockers hit (all real, all fixed)
1. **`GatherQMM::vjp` error** — MoE routing indices come from `argpartition`
   (non-differentiable) but the quantized gather-matmul VJP tried to differentiate
   them. Fix: `mx.stop_gradient(indices)` at the `SwitchGLU` boundary.
2. **Recompilation every step** — variable prompt lengths → new graph each step
   (45–80 s/step). Fix: pad all to `max_tokens`, mask the loss.
3. **Adam optimizer state OOM** — 5.3 B trainable scale params × 2 momentum buffers
   ≈ +59 GB. Fix: `--scales-only` + `--optimizer sgd` → stable 109 GB, ~45 s/step.
4. **Gradient accumulation made things WORSE** — every variant (fp32, +gc/clear_cache,
   fp16) *raised* memory (115–130 GB) and tripled step time. Abandoned.
5. **Divergence at lr=1e-4** — loss climbed 2–5× above the untrained baseline and did
   not recover. Diagnosed with `diag_specific_kl.py` (trained loss vs untouched-student
   baseline on the *same* prompts). Fixed with lr=1e-6, then 1e-7.

### Why its value is questionable
The **untrained** student's baseline KL vs the teacher was already **0.06–1.4** — low.
There is little gap left for DWQ to close, and it was never shown to beat noise.
Effort is better spent on calibration quality (§2, §3) than DWQ tuning.

---

## 8. Quantization format

| Format | scales | biases | DWQ-trainable? | group_size |
|---|---|---|---|---|
| **affine 4-bit** | float32 | float32 | ✅ yes | 64 |
| **nvfp4** | **uint8** (FP8 microscales) | none | ❌ no | **must be 16** |

- nvfp4 is the better raw 4-bit format but stock DWQ cannot train its uint8 scales.
- **nvfp4 requires no calibration data at all**, which given §3 is a real advantage.
- Vision tower left **unquantized (BF16)** — the deploy model keeps full vision.
- LM Studio runs `step3p7` + affine; **DWQ is invisible to the runtime**.

Built student: `models/Step-3.7-p15-4bit` — 92 GB, 245 experts, **4.632 bpw**,
via `scripts/build_student.py` (fused apply+quantize, never writes the ~316 GB
reaped-BF16 intermediate). Smoke-tested: loads, generates coherently, 58 tok/s.
**Coherent with zero recovery**, corroborating the 15% choice.

---

## 9. Router architecture (relevant to proposed "Router Bias Surgery")

```python
corrected_scores = scores + router_bias          # SELECTION only
topk_indices = argpartition(-corrected_scores, kth=top_k-1)[..., :top_k]
topk_weights = take_along_axis(scores, topk_indices)   # weights use UNBIASED scores
topk_weights = topk_weights / sum(topk_weights)        # norm_expert_weight: True
```

Three consequences:
1. **No routing void.** `argpartition` runs over the **245 survivors**, so every token
   still gets exactly 8 experts. A rank-9 kept expert is promoted automatically.
2. **Mass is already conserved** by `norm_expert_weight: True`.
3. **Relative ordering among survivors is unchanged** — each keeps its original bias.

`router_bias` is a DeepSeek-style aux-loss-free **load-balancing** term; load imbalance
costs training throughput, not inference quality (no capacity limits at inference).
Proposals to "absorb probability mass" via bias adjustment therefore address a problem
this architecture already handles.

---

## 10. Untested proposals — status

| Idea | Assessment |
|---|---|
| **REAM** (merging instead of pruning) | Genuinely well-matched: this model *is* prune-resistant, and REAM consumes REAP scores so existing saliency is reusable. Biggest build; unvalidated here. |
| **DPP-REAP** (diversity selection) | Math is correct. But the co-occurrence kernel is **~36× under-sampled off-diagonal** (2.8% vs 0.077% of tokens), so dividing by `total_tokens` makes it diagonally dominant → **degenerates back to REAP**. Fix: normalize by co-occurrence count. Naive greedy is also ~2.6×10¹¹ flops/layer; incremental Cholesky reduces it to ~1.7×10⁷. |
| **Router Bias Surgery** | Addresses a non-problem — see §9. |
| **RMSNorm residual projection saliency** | The strongest untested idea. RMSNorm's Jacobian contains `P⊥ₓ = I − xxᵀ/‖x‖²`, so expert output parallel to the residual stream is rescaled away and should not count toward saliency. Cheap, drop-in, independent. |
| **Layer-adaptive allocation** | Already measured: **5.3%** at p15 (§5). |

### Cheapest decisive next experiment
Add co-occurrence tracking to `LayerSaliency` and measure **off-diagonal mass** of the
correlation matrix. If `C_{i,j} ≈ 0`, experts are already orthogonal and **both DPP and
REAM are ruled out** for ~40 min of machine time.

**Note:** co-occurrence is inherently *pairwise*; everything currently stored
(`reap`, `freq`, `gate_sum`, `reap_count`) is a **per-expert marginal**. It cannot be
recovered from existing artifacts — a new pass is required. But **5k is unnecessary**:
at 1000 prompts each pair averages ~694 co-occurrences (threshold is 5), ample for a
yes/no redundancy gate. ~40 min vs ~3 h.

---

## 10b. Perplexity evaluation — two methodology bugs found

Building the first *non-proxy* measurement (`reap_stream/eval_ppl_streamed.py`)
surfaced two bugs that would have made every PPL number meaningless. Both were
caught by noticing an implausible result, not by the code failing.

Initial smoke test on 6 held-out prompts gave **agentic PPL = 388** — absurd for a
frontier model, while `reasoning_math` scored a healthy 3.30. That split was the clue.

| Category | Original | + `head` trunc | + raw text (**correct**) |
|---|---|---|---|
| agentic | 388.41 | 220.56 | **10.75** |
| coding | 42.30 | 42.30 | **9.13** |
| reasoning_math | 3.30 | 3.17 | **2.55** |
| **OVERALL** | 42.94 | 33.22 | **6.13** |

**Bug 1 — `headtail` truncation is wrong for perplexity.**
It's correct for *saliency* (sampling both task setup and answer gives representative
routing statistics), but for perplexity it splices first-512 onto last-512, creating a
hard discontinuity mid-sequence. Tokens after that seam are genuinely unpredictable,
inflating NLL for reasons unrelated to model quality. Agentic prompts (5000–10000
tokens) are the most truncated, hence worst affected; `reasoning_math` mostly fits
under 1024 and was untouched. **PPL eval now defaults to `head`** (contiguous prefix).

**Bug 2 — chat-template double-wrapping** (inherited from `_tokenize_prompts`).
The calib rows already carry their own `SYSTEM:/USER:/ASSISTANT:` structure, but
`apply_chat_template` wrapped the whole thing *again* as a single user turn:
```
<|begin_of_sentence|><|im_start|>user\nSYSTEM:\nYou are a helpful assistant...
```
The model receives `SYSTEM:` as literal text inside a user message — a structure it
never saw in training. This is why agentic (elaborate embedded SYSTEM blocks + function
schemas) and coding suffered most, while plain `USER:/ASSISTANT:` reasoning_math barely
moved. Fixing it gave a **36× improvement on agentic**. PPL eval now feeds raw text;
`--chat-template` opts back in.

> **Does this invalidate the saliency work? No.** Every saliency run used identical
> treatment, so comparisons *between* them (old vs new plan, text vs vision) remain
> valid — it's a level shift, not a differential one. It would only have corrupted
> perplexity, which is exactly where it was caught.

**Eval protocol:** 500 held-out prompts from rows 5000+ (never used in calibration —
saliency used the first 2500, DWQ the first 1500), multimodal excluded (misaligned and
imageless, unanswerable as text), `head` truncation at 1024 tokens, raw text, with
**per-category breakdown** — an aggregate number can hide "coding fine, agentic
degraded 5%".

Measured BF16 streaming cost: **0.86 s fixed/prompt + 1.13 ms/token** → ~17 min for
500 prompts. Resident quantized models are far faster (~1–2 min, no per-layer disk reads).

---

## 11. Method lessons

- **`footprint <pid>`, not MLX counters.** MLX under-reported by >2×.
- **Compressor + swap are the crash predictors**, not free RAM.
- **MLX hoards freed buffers** — `mx.clear_cache()` inside hot loops, not just at phase
  boundaries. Single highest-leverage memory fix found.
- **List comprehensions over big tensors double peak memory** — update in place.
- **Never compare against a partial checkpoint.** Zero-filled layers silently produce
  meaningless correlations that look like real structure (§2 correction).
- **Pair multimodal data in one pass.** Index-based pairing across separate loads
  desyncs silently.
- **Test the fix, not just the problem.** `tail` truncation and gradient accumulation
  both sounded right and both measured *worse*. Every fix that survived here was
  validated against ground truth first.
- **An implausible number is a bug report.** Agentic PPL of 388 wasn't a finding about
  the model — it was two stacked methodology bugs (§10b). Sanity-check magnitudes
  against what the quantity *should* look like before interpreting it.
- **The right setting is task-dependent.** `headtail` is correct for saliency and wrong
  for perplexity; the same knob flips depending on whether you're sampling routing
  statistics or modelling a contiguous sequence.
- **Most numbers here are still proxies.** Perplexity (§10b) is the first non-proxy
  measurement. Real task benchmarks (SWE-bench, Terminal-Bench — what this model is
  actually built for) remain unrun.