[grug] FSDP hero optimization sweep: protocol and results

Tracking issue for testing the optimizations identified in #8053 against the 1-rack FSDP hero. Analysis and profile links are in #8053; this issue records the measurement protocol and the measured result of each arm. Measurement protocol Each arm is a 20-step single-rack run (64x GB200, batch 1024 x seq 4096) with `--no-save-checkpoints --watch-interval 0`. Steps 0-1 are compile and the PGLE recompile (146 s and 108 s in the reference run); steps 2-4 absorb the one-time first-batch data-loader stall. The score is the median `throughput/duration` over steps 5-19. Median rather than mean, because the steady-state distribution is right-skewed. The reference run (`gprof-xprof-1rack-20260807`) over steps 5-19 gives median 18.086 s, MAD 0.077 s, min 17.936 s, max 18.699 s, stdev 0.210 s (1.16%). The median's standard error at n=15 is about 0.068 s, so a 2-sigma within-run effect is 0.14 s, roughly 0.15 MFU points. Within-run noise is not the binding constraint; run-to-run variation from node allocation and recompilation is. Every wave therefore launches a `control` arm with the unmodified configuration in the same wall-clock window, on one of the four racks. An arm counts as a win only when it beats the control by more than the control-to-control spread. MFU is reported for readability but carries no information beyond step time, since the FLOP denominator is fixed by the shape. Four racks at production priority run one control plus three experiment arms per wave, about 15 minutes per wave. Arms to test, in expected-uplift order | Wave | Arm | Change | Expected | |---|---|---|---| | 1 | `simple` | `NCCL_PROTO=Simple` | +1.5 to +3 | | 1 | `ll128` | `NCCL_PROTO=LL128` | +1.5 to +3 | | 1 | `nvls` | `NCCL_ALGO=NVLS,Ring`, `NCCL_NVLS_ENABLE=1` | +1.5 to +3 | | 2 | `combine` | Raise the XLA all-gather and reduce-scatter combiner thresholds | +0.3 to +1 | | 2 | `pipelined` | `--xla_gpu_enable_pipelined_all_gather`, `..._reduce_scatter` | +0.3 to +0.8 | | 2 | `nopgle` | `JAX_ENABLE_PGLE=0`, isolated from the nsys overhead that confounded the first estimate | 0 | | 3 | `shard-small` | FSDP-shard router, `attn_gate`, and both `GatedNorm` factor pairs | +0.5 to +0.8 | | 3 | `interleave` | Interleave w13 before the expert-weight all-gather instead of after | +0.6 to +1.0 | | 3 | `chunks` | `expert_chunks` sweep and an unequal first chunk | +0.3 to +0.7 | | 4 | `offload-act` | `remat_mode="save_moe"` with `save_and_offload_only_these_names` to Grace | +2 to +4 | | 4 | `wgrad` | MoE weight-grad on the CUTLASS grouped GEMM instead of XLA `ragged_dot` | +1 to +1.5 | | 5 | `combined` | Every arm that won, stacked | — | Peak HBM logging lands before wave 3 so every later arm reports `peak_bytes_in_use`; today nothing in the training path records it, which is why the 160 GB to 154 GB regression could not be attributed from the profile. Results are posted per wave as comments, and consolidated into this issue body at the end. Part of #8053