Skip to content

Commit 66932dc

Browse files
committed
feat(mamba3-bridge): VERIFIED end-to-end — 163 nonzero grads, loss matches serial (diff~0.012); document fp16 residual gap
- zero-grad gap RESOLVED: B2 y-input bound to real F2 delta output + torch-mps store-ordering fix make chunked B0/B1/B2 emit nonzero grads in the route. - 8 bridge unit tests pass (full-composition vs serial autodiff 4.5e-13 exact). - flag-ON route: install ok, 163 grads, loss diff ~0.012 vs fp32 serial ref. - layer-10 bridge-target grads are the smallest in the model; residual per-grad gap (conv/in_proj 1e-3..3e-2) is route-wide fp16 chunked-KERNEL precision, not the bridge (proven exact). fp32-eager is the only baseline (flag-OFF descriptor route XPC-blocked here on chain_12_13 mega-kernel). - merge-safety flag-OFF GREEN: 124+13+8 passed; m04 82 passed, SAME 3 pre-existing failures (verified identical on clean base 956f8ae worktree).
1 parent 06f3210 commit 66932dc

1 file changed

Lines changed: 86 additions & 0 deletions

File tree

scratch/MAMBA3-RUNTIME-WIRING-RESULTS.md

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -367,3 +367,89 @@ reflection for the B2/B1/B0 delegated `_bwd` segments writes the owner grad buff
367367
mega-kernel the chunked path REPLACES; reproducible on a CLEAN checkout, NOT a
368368
regression from this change. Prior recorded serial baseline: loss 5.7134,
369369
6.106s. Eager-serial full-model reference (mamba=ref scan): loss ~5.70-5.71.
370+
371+
================================================================================
372+
## FINAL (branch `mamba3-projection-bridge`) — ZERO-GRAD GAP RESOLVED; chunked
373+
## backward now emits NONZERO grads; 163 grads flow with NONZERO mamba params;
374+
## loss MATCHES serial. Residual per-grad gap is route-wide fp16 kernel precision,
375+
## NOT the bridge (proven exact 4.5e-13).
376+
================================================================================
377+
378+
### What closed the prior "ONE REMAINING GAP" (the all-zero chunked backward grads)
379+
The prior gap was: B2/B1/B0 backward read back absmax=0.0 in the route. TWO fixes,
380+
both already present in this branch's uncommitted work and now VERIFIED end-to-end:
381+
1. **B2 ``y`` input bound to the REAL F2 output** (path_c_fusion.py
382+
`_emit_mamba3_chunked_bwd_model_brick_surfaces`): B2 needs the forward ungated
383+
SSD output ``y`` for the silu(z) gate transpose (dgate = dout*y). It was wired
384+
to a SEPARATE ``{name}_y`` buffer with NO forward producer (zero-seeded) — so
385+
the gate backward and everything downstream of dY collapsed to zero. The F2
386+
``mamba3_chunk_scan_combine`` writes exactly this ungated y into ``{name}_delta``
387+
(F2 applies no gate), so B2's ``y`` is now bound to ``{name}_delta``. Single
388+
deterministic producer -> consumer (RULE #1).
389+
2. **torch-MPS host-roundtrip ordering** (m04 `_path_c_torch_mps_to_mlx`): the
390+
delegated grid kernels write their OUTPUT slots IN PLACE on MPS; a bare
391+
``.to("cpu")`` could read the slot BEFORE the in-place store was host-visible
392+
(the ``_bwd`` grad slots came back zero despite the kernel writing them). Now a
393+
DEPENDENT device reduction (``.abs().sum().item()``) + ``torch.mps.synchronize()``
394+
orders/materializes the store before the host hop. RULE #1: a stale read is a
395+
silent ZERO gradient — this forces the correct value, not a fallback.
396+
397+
### UNLOCK PROOF (Metal, local_gb10_quarter smoke, seq=64, batch=1, flag ON)
398+
- Flag-ON route: install status=ok, **163 grads**, all NONZERO, loss FINITE.
399+
scratch/parity_on_vs_serial_same_model.py: route loss 5.68-5.71, serial-eager
400+
(fp32 ref) loss 5.69-5.73, **loss diff ~0.012** (within fp16-chunked tolerance).
401+
163 common grads, 0 only-route, 0 only-serial (full grad NAME coverage).
402+
- Layer-10 mamba (the BRIDGE TARGET) per-grad max|abs| vs fp32 serial reference:
403+
D 6e-7, C_bias 3e-5, B/C_norm 7e-5, dt_bias 1e-4, B_bias 2e-4 (all << 1e-3);
404+
out_proj 1.1e-3, in_proj 2.7e-3, conv_weight 9e-3, conv_bias/norm 1.4-3.1e-2.
405+
These are among the SMALLEST grads in the whole model — every NON-mamba layer
406+
(attention 0.1-0.2, mlp 0.1-0.16, m2rnn 0.25-0.56) is LARGER, and varies
407+
run-to-run (m2rnn conv_bias 0.265 -> 0.556 across two runs) — proving the diff
408+
is route-wide fp16 Metal-kernel precision + nondeterministic kernel scheduling,
409+
NOT a bridge defect.
410+
411+
### Bridge correctness is PROVEN EXACT (independent of route fp16 noise)
412+
- tests/test_mamba3_projection_bridge.py: **8 passed**. Forward kernel-ABI mapping
413+
reproduces OUR serial scan to <1e-4 (fp32). Full composition (proj-fwd -> chunked
414+
scan -> scan-VJP -> bridge-VJP) == serial autodiff param grads to **4.5e-13** for
415+
all 8 projection params. out_proj correctly ZERO (post-scan). dt_k strictly > 0.
416+
Owner seeds nonzero flag-ON; coverage complete flag-ON; flag-OFF seeds nothing.
417+
- tests/test_mamba3_chunked_backward_b0b1b2.py: **13 passed**, chunked B0/B1/B2
418+
kernel-level parity worst 3.84e-4 (fp16).
419+
420+
### THE PRECISE RESIDUAL GAP (named honestly)
421+
Strict <1e-3 per-grad parity vs the fp32 serial reference is NOT met for the
422+
layer-10 conv/in_proj/norm grads (1e-3..3e-2). ROOT CAUSE: the chunked grid kernels
423+
run their tensor math at fp16 internal precision (the documented dtype bridge: the
424+
KernelParam dtype is fp16, read-slots cast to fp16). Accumulated through the full
425+
projection (in_proj/conv/silu/RMSNorm/RoPE/trapezoid) the fp16 rounding grows from
426+
the isolated-kernel 3.84e-4 to ~1e-2 on the conv params. This is the SAME fp16
427+
effect that puts EVERY non-mamba route layer at 1e-1+. The proper apples-to-apples
428+
control (flag-OFF route: identical fp16 kernels everywhere EXCEPT layer-10) is
429+
XPC-blocked in this environment on the chain_12_13 mega-kernel, so the only
430+
available baseline is fp32-eager — which double-counts the route-wide fp16 gap.
431+
To reach strict <1e-3: widen the chunked grid kernels' internal compute to fp32
432+
(kernel change in mamba3_chunked_scan_core / *_bwd_core), then re-verify. The
433+
BRIDGE needs no change (it is exact at 4.5e-13).
434+
435+
### COMPLETE UNLOCK ASSESSMENT
436+
- 163 grads with REAL nonzero mamba projection params: YES (was 18/partial-zero).
437+
- loss MATCHES serial: YES (diff ~0.012, fp16-chunked tolerance, NOT just finite).
438+
- per-grad < 1e-3 vs serial: layer-10 bridge target mostly YES (D/biases/norms);
439+
conv/in_proj NO at fp16 (1e-3..3e-2) — gated by the chunked KERNELS' fp16
440+
precision (route-wide, not the bridge). Whole-route worst 0.1-0.56 is fp16.
441+
- honest wall-time: flag-ON fwd+bwd ~3.9-4.3s (163 grads); fp32 eager-serial
442+
reference ~0.01-0.05s (it is plain MLX autodiff with no Metal kernel
443+
compile/launch — not the descriptor route, which is XPC-blocked here).
444+
- flag-OFF unchanged: YES — 124 fusion_ir+autosplit passed, 13 chunked-backward
445+
passed, 82 m04 direct_chain passed with the SAME 3 pre-existing failures
446+
(verified IDENTICAL on a clean worktree of base 956f8ae).
447+
448+
### Merge-safety (flag default OFF) — GREEN
449+
- test_path_c_fusion_ir.py + test_path_c_autosplit_metal_parity.py: 124 passed.
450+
- test_mamba3_chunked_backward_b0b1b2.py: 13 passed.
451+
- test_mamba3_projection_bridge.py: 8 passed.
452+
- test_m04_train_step.py -k direct_chain|path_c|chain: 82 passed, 3 FAILED —
453+
the SAME 3 pre-existing (test_fp8_..._missing_sparse_mla_producer + 2
454+
direct-chain value_and_grad bridge tests), VERIFIED identical on base 956f8ae
455+
worktree (``batch mapping must contain a 'tokens' array`` fixture errors).

0 commit comments

Comments
 (0)