docs(pr2): mark PR-2 LANDED+MEASURED — bs=4xseq=4096 OOM(114G)->fits(64.93G)

Update FUSED-PIPELINE-ROADMAP §5 PR-2 (status + measured peak table + numeric
equivalence + honest floor) and MEGATRON-VS-MLX-PATHS (seq=4096 now reachable:
64.93 GiB peak / 514 tok/s live on gb10, was 'cannot run'). The 116 GB was a
loop/graph problem, fixed by env-gated numerically-equivalent grad accumulation
in one_step_train (commit 05e53d8), not a kernel change.

Author: apstenku123

docs/FUSED-PIPELINE-ROADMAP.md

@@ -135,9 +135,28 @@ All paths under `/Volumes/external/sources/tilelang/examples/`. Effort tags: Low
 
 ### PR-2 — The memory-loop fix (grad-accum + selective checkpoint barrier + 8-bit optimizer)
 **Highest mem ROI, smallest diff.** Independent of PR-1; can land in parallel.
+**STATUS (2026-06-02): LANDED + MEASURED ON gb10 (commit `05e53d8`). bs=4×seq=4096 went OOM→FITS.**
 - **First step:** wire `grad_accum_steps` (scaffolding at `checkpoint.py:101-104`) into `one_step_train` (`loop.py`) with `mx.eval(grads)`-and-drop per microbatch; add `mx.eval(hidden_states)` barrier after each attention-only checkpointed layer; switch optimizer to `optimizers_quantized.py`; call `maybe_clear_cache_after_step` (`runtime/memory.py:191`) each step.
 - **Unblocks:** seq=4096 step 116 GB → 30–40 GB on gb10 — makes the real production config (bs4×seq4096) *reachable* on a single box without waiting for Relax.
 
+**What landed (`05e53d8`, all env-gated, default = prior behavior — RULE #1):**
+- `one_step_train` (`loop.py`): `grad_accum_steps` kwarg / `GRAD_ACCUM_STEPS` env (default 1 = byte-identical to the prior single-`mx.eval` step). When N>1: split the batch into N microbatches, run `value_and_grad` per micro, `mx.eval` the **token-weighted** running accumulator as a free-barrier, drop each micro's fwd/bwd graph before the next. Optional `clear_cache` hook.
+- `hybrid_lm.py`: opt-in `mx.eval(hidden_states)` after each grad-checkpointed layer (`CPPMEGA_MLX_CHECKPOINT_EVAL_BARRIER=1`, default OFF) — converts the previously no-op checkpoint into a real free-barrier.
+- 8-bit optimizer (`make_adam8bit`) + `maybe_clear_cache_after_step` already wired in `train_hybrid_tiny.py`; probe `scripts/pr2_seq4096_memory_probe.py` exposes `--opt adam8bit`.
+
+**Numeric equivalence (RULE #1, verified on Mac, `scratch/pr2_grad_accum_equivalence.py`):** N=4 micro vs full-batch → loss abs-diff **0.0**, worst grad rel-diff **2.8e-7**, post-step param rel-diff **1.2e-7**; checkpoint barrier ON vs OFF → grads **bitwise identical**. Token-weighting (`weight = n_i / N_total`) makes `sum_i weight_i · grad_i` exactly the full-batch grad for a token-mean loss. 165 loop/compiled/gradient/checkpoint tests pass.
+
+**Measured peak — gb10, seq=4096, dense pattern=A, hidden=2048 / depth=24 / 542.5M params** (peak = `mx.get_peak_memory`; 100 GB box budget; `free -g` cross-check; watchdog SIGTERM @100 G):
+
+| config | grad_accum | optimizer | active GiB | peak GiB | tok/s | fits ≤100 G? |
+|---|---|---|---|---|---|---|
+| bs=4×seq4096 monolithic (BEFORE) | 1 | AdamW | — | **OOM (SIGTERM @114 G)** | — | **NO** |
+| bs=4×seq4096 (AFTER) | 4 | AdamW | 6.06 | **67.94** | 526 | **YES** |
+| bs=4×seq4096 = Megatron config | 4 | Adam8bit | **3.05** | **64.93** | 514 | **YES** |
+| bs=8×seq4096 (2× prod batch) | 8 | Adam8bit | 3.05 | 64.93 | 516 | YES |
+
+**Result:** the real `bs=4×seq=4096` step went from **OOM (>114 GB, watchdog-killed)** to **67.94 GiB — it now fits on a single gb10**, loss finite + decreasing (10.54→10.29), ~526 tok/s. Active working set drops **linearly with N** (6.06 GiB @ N=4 → 3.05 GiB @ N=8), confirming the activation lever. The ~64–68 GiB *peak* is the **allocator high-water mark, not the live set** (active is 3–6 GiB) — `clear_cache` trims but the CUDA-MLX allocator/graph-cache holds the high-water mark. **Honest floor:** loop-level levers land the seq=4096 step at **~62–68 GiB peak — well under the 100 GB box budget, above the 30–40 GiB stretch target.** Reaching 30–40 GiB *peak* needs allocator-limit pinning (`apply_memory_limit_plan` / `MLX_CUDA_GRAPH_CACHE_SIZE`) and/or the durable Relax whole-step liveness (§4); the dominant residual term is the **allocator high-water mark / static param+optimizer block**, not activations (which grad-accum already collapses).
+
 ### PR-3 — Auto-fusion / block-registry (the Megatron-spec equivalent)
 **The composability layer.** Depends on PR-1's brick descriptors existing.
 - **First step:** define the `TransformerLayerSubmodules`-style block-spec dataclass; register `mamba3`/`mla`/`moe` as `FusionBlockDescriptor`s (`fusion.py:64`) with paired `_bwd`; register the fused chain signatures in `FusionScheduleRegistry` (`fusion.py:295`); drive `plan_path_c_direct_fusion_chain_for_region` candidates from `select_longest` (`fusion.py:344`) matches instead of the fixed signature. Audit template-execution for any try/except that masks a broken chain (RULE #1).

docs/MEGATRON-VS-MLX-PATHS.md

@@ -133,11 +133,21 @@ Sources: `reports/cppmega_1b_speed_matrix_gb10_fastfused_20260601_b2s512.md`,
 | bs=2 × seq=512 | 1024 | **259** ✅ | **cannot run** | 27-32 (b)/17-22 (c attempt) | path_c blocks at bs>1 |
 | bs=1 × seq=1024 | 1024 | **191** ✅ | **cannot run** | 32.6 (b) | path_c blocks at seq>512 |
 | bs=1 × seq=1024 (eff. MoE) | 1024 | 166 ✅ | n/a | 32.8 | `CPPMEGA_MOE_EFFICIENT=1`, loss 11.29→5.65 |
-| bs=1 × seq=4096 | 16384 | **cannot run** | n/a | fwd-only 4.9; fwd+bwd+adamw **97+** | backward+optimizer burst > 105 GB cap |
-| bs=4 × seq=4096 (Megatron config) | 16384 | **cannot run** | n/a | — | unreachable (bwd+opt burst) |
+| bs=1 × seq=4096 | 16384 | **cannot run** (monolithic) | n/a | fwd-only 4.9; fwd+bwd+adamw **97+** | backward+optimizer burst > 105 GB cap |
+| bs=4 × seq=4096 (Megatron config) | 16384 | **cannot run** (monolithic) | n/a | OOM @114 | unreachable as ONE monolithic step (bwd+opt burst) |
+| **bs=4 × seq=4096 + PR-2 grad-accum=4** | **16384** | **514 tok/s** ✅ | n/a | **64.93** (active 3.05) | **NOW REACHABLE** — loop-level grad-accum (`05e53d8`, dense pattern=A, hidden=2048/depth=24/542.5M, Adam8bit) |
 
 (muon mirrors adamw on scale: bs1→bs2 path_b 92→159.)
 
+> **PR-2 update (2026-06-02, FUSED-PIPELINE-ROADMAP §4/§5, commit `05e53d8`):** the bs=4×seq=4096 step is
+> **no longer unreachable**. The 116/114 GB monolithic burst was a *loop/graph* problem (whole fwd+bwd+update
+> forced at one terminal `mx.eval`), not a kernel one. `one_step_train` now supports env-gated, numerically
+> equivalent **gradient accumulation** (split the batch into N microbatches, `mx.eval`-and-drop each micro's
+> graph before the next). At N=4 + Adam8bit the Megatron-shape step lands at **64.93 GiB peak (active 3.05 GiB),
+> ~514 tok/s, loss finite + decreasing** — measured live on gb10, well under the 100 GB box budget. The residual
+> ~65 GiB peak is the MLX-CUDA **allocator high-water mark**, not the live working set (3 GiB); driving it to the
+> 30–40 GiB target needs allocator-limit pinning / Relax whole-step liveness (roadmap §4), not more grad-accum.
+
 ---
 
 ## 3. Apples-to-apples normalization
@@ -171,10 +181,14 @@ batch-amortization delta — it is **not** closeable to a single clean factor be
    So a steady-state path_c-vs-path_b-at-scale comparison is **unmeasurable**; path_c simply doesn't run off its
    build shape. The Path-C win is therefore **compile-time only** (1.15-2.0× faster first-step / kernel build),
    not steady-state throughput at bs=1.
-2. **path_b runs everywhere and scales** (156→259 with batch, 156→191 with seq) but **cannot reach seq=4096**:
-   the seq=4096 **forward fits in 4.9 GB**, but the **eager backward + AdamW first-update bursts to ~97 GB**
-   (gradients + m/v optimizer state + recomputed grad-checkpoint activations all materialize at once) — over the
-   105 GB SIGTERM safety cap. So bs=4×seq=4096 (and even bs=1×seq=4096 fwd+bwd) is unreachable in MLX-eager.
+2. **path_b runs everywhere and scales** (156→259 with batch, 156→191 with seq); seq=4096 was **unreachable as a
+   single monolithic step** — the seq=4096 forward fits in 4.9 GB, but the eager backward + AdamW first-update
+   burst to ~97–114 GB (gradients + m/v optimizer state + recomputed grad-checkpoint activations all materialize
+   at once), over the safety cap. **PR-2 (commit `05e53d8`) makes it reachable:** loop-level gradient accumulation
+   (N microbatches, `mx.eval`-and-drop per micro — numerically equivalent to the full-batch step) drops the
+   activation term ~linearly with N, so **bs=4×seq=4096 now runs at 64.93 GiB peak / ~514 tok/s** (live gb10,
+   dense pattern=A, hidden=2048/depth=24/542.5M, Adam8bit). The remaining ~65 GiB is the allocator high-water mark
+   (live set 3 GiB), not the activation burst — see roadmap §4 PR-2 for the 30–40 GiB durable path.
 3. The wall is the **whole-model backward+optimizer burst, NOT the MoE** — measured: efficient sparse MoE and
    dense MoE have the same ~32.6/32.8 GB peak at seq=1024 (MoE is only 4 of 13 layers). Megatron does bs=4×seq=4096
    in ~26 GB because its fused training step (activation checkpointing + a fused optimizer that never