Skip to content

⚡ Phase 16 — Training throughput

WP-177 — one MuSGD launch per exact matrix shape

The measured cost was scheduling. The n-scale detector has 127 matrix parameters but only 41 exact (shape, dtype, device) keys — a 3.10x reduction in launches rather than the 6.35x a census of 20 would have promised, and the ratio rises with scale (s 3.10, m 4.23, l 5.74, x 5.74) — so _step_group now stacks equal-key matrices inside their existing optimizer group and maps the existing _parameter_update over each stack. The map is taken on CUDA alone, and a key holding one parameter is called directly on every backend. _parameter_update, _muon_branch and orthogonalize still define the result; parameter groups, shapes, dtypes and devices never mix, while vectors and w_muon == 0 keep their existing path. On an NVIDIA L4 (driver 580.173.02, torch 2.9.1+cu129, CUDA 12.9), the real Lightning trainer and loader ran n-scale detection over 128 COCO images at batch 4, 640 px, BF16 and two workers. Across five fresh ABBA processes after warm-up, median step interval moved from 333.56 ms (331.23--338.37 ms) to 208.29 ms (206.83--210.54 ms), a 37.11% reduction; every repeat improved by 36.81--38.87% and peak allocation stayed 809 MiB. The earlier 40.87% prototype duplicated the expression in a benchmark subclass and is superseded by this production-patch result.

The formula stays fixed; the bit schedule does not. This change retires _parameter_update's previous promise that a w_muon > 0 run was bit-for-bit unchanged. On an L4 with torch 2.14.0+cu130, direct execution and a singleton vmap are bit-identical across all 20 real matrix shapes, showing that the CUDA difference begins when equal shapes are actually batched; on CPU with the same torch build, 19/20 are identical and (4, 21, 1, 1) differs by at most 4.76837158203125e-07, so no backend-wide bit guarantee replaces the old one. With the real bucket sizes and fixed gradients, the maximum parameter difference from the old loop is 3.814697e-06 after 60 local CUDA steps and 1.192093e-06 after 35 CPU steps. make gate is green with 3,260 passed, 97.31% coverage and 50/50 offline goldens; make gate-gpu passes all six accelerator/data tests but its golden sweep is 54/55, with keypoint train OKS AP 0.691134 against 0.574728 +/- 0.05. A same-L4 ABBA control gives baseline 0.5315, 0.5315 and this row 0.6911, 0.6911. The higher value clears the task floor but proves no long-run quality gain; no golden moves, WP-177 stays open, and the review owns whether this deterministic trajectory change is acceptable.