Refactor the MXFP8 fused-MLP overrides onto stock parameters

[wolfcomos]

Review vehicle for the self-owned MXFP8 fused-MLP refactor. Base = swiglu-mxfp8-upstream, so this diff is exactly what would land on the open upstream PR (pytorch#4257) on approval. Two commits (honest trail of the design evolution; can squash at approval):

1. 88e2e8928 — rebase the grouped class off FusedGroupedExperts onto stock GroupedExperts (self-owned w13), umbrella naming (mxfp8_fused_mlp.py, MXFP8FusedMLP/MXFP8FusedGroupedMLP, one test file).
2. 9f44233d3 — drop the fused w13 parameterization entirely, dense included: both overrides keep STOCK parameters (w1/w2/w3 Linears; w1_EFD/w2_EDF/w3_EFD) and pack the two projections at forward time (torch.stack(..., dim=1/2) — byte-identical to the old merge-hook mapping). All parameter-lifecycle machinery deletes: no param surgery, no state-dict hooks, no init/sharding remaps. The module ends with zero fused_swiglu.py imports; overrides/fused_swiglu.py (merged-in-main) stays byte-untouched. A future fully-fused grouped path wires in as just another forward-pack + composite (3-line docstring note is the only structural gesture — no knob shipped here).

Why stock params

• Checkpoints AND fresh-init are bitwise-identical to stock modules by construction (no hooks needed; new CPU test compares seeded fresh-init state-dicts against genuine stock builds, torch.equal per key).
• Any future fusion plan (e.g. the grouped-GEMM kernels' 32-block layout) differs only in its pack function — no plan-dependent parameter layout, hooks, or checkpoint divergence.
• Cost: one bf16 weight stack per forward.

Size

Cumulative diff vs main: +846 / −0 (4 files: module 619, tests 200, two registry flavors 27). No validation check or test assertion removed; suite is 8 tests in one file.

Evidence (raw logs in agent_scratch/upstream_swiglu/validation/{fused_mlp_refactor,stock_params}/)

• Grouped path, strict checkpoint-paired bitwise vs Add MXFP8 fused-MLP overrides in torchtitan pytorch/torchtitan#4257's head (shared seed checkpoint, NGPU=2 EP=2, 5 deterministic steps): every printed loss and grad-norm equal.
• Dense path: grad-norm column equal all 5 steps; loss differs in the last printed digit at steps 3 and 5 under active grad-clipping — root-caused and ablated: with clipping neutralized (max_norm=1e8, clip coef exactly 1.0) both arms are exactly equal all steps. The sole divergence is the clip total-norm's FP reduction shape (one (H,2,D) grad vs two (H,D) grads), inherent to dropping the fused parameter, not a packing defect.
• Fresh-init state-dicts bitwise-equal to stock modules (both paths, seeded CPU test).
• Suite 8/8 on GPU, 8/8 GPU-less, and 8/8 executed (not skipped) on isolated stock PyPI torchao 0.18.0 with no GPU — the upstream-CI condition.
• flake8/py_compile clean; overrides/fused_swiglu.py byte-identical to main.

On approval: fast-forward both commits onto swiglu-mxfp8-upstream (no history rewrite) and refresh the pytorch#4257 title/description.

🤖 Generated with Claude Code