# 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 109 GB IOAccelerator (graphics) [all DIRTY, 0 reclaimable] phys_footprint: 110 GB ``` **Always use `footprint `, 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 `, 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.