📦 WP-193b — perf(optim): batch compatible MuSGD matrices and momentum¶
Engineering log · Phase 17 — R1 pose-objective parity
🔗 Depends on WP-193 · Needed by none
WP-193b — compatible matrix batches and momentum dispatch¶
Compatible matrix batches and momentum dispatch. One diagnostic session after the native-batch and image-cache changes attributed a 49.28 ms median to the optimizer within a 159.60 ms complete detection interval on an isolated L4. Its 10.90 ms momentum-dispatch median suggested a 6.8% stage share; this is an exploratory opportunity estimate, not an independently repeated speedup or a universal ceiling. A single bounded-padding-only complete-path screen suggested 5.94%, without a run-level range; it did not establish a standalone improvement above the investigation's 8% bar. The combined candidate targets the remaining launch overhead within the same optimizer phase.
Zero extension changes the batch, not the quintic. For a wide matrix X, appending zero columns preserves its Frobenius norm and X @ X.T; every iterate of the existing polynomial therefore keeps those columns zero. CUDA matrices sharing the smaller dimension can run in one native batch, then crop and transpose back before their original dimension-dependent scale, gains and decay. Each padded gradient stack is bounded by the parameter and gradient bytes of the matrices in its own batch. A singleton, or a batch whose matrices all share one shape, keeps the exact-shape update; equal-shape matrices that share a smaller dimension with a matrix of another width join the padded batch. Parameters whose memory ranges overlap keep exact-shape keys, whichever storage object wraps them. No coefficient, iteration count or normalization rule changes.
Momentum keeps its rounded operations. Three foreach operations replace the CUDA per-parameter multiply, add and Nesterov add, with the original operation order. Device and dtype buckets keep independent state. Shared parameter/gradient/buffer storage and unusual restored state use the original scalar method; CPU and MPS keep their previous implementation. No public flag or checkpoint field is added.
The numerical contract is bounded agreement. Mixed-width GEMMs and reductions may round differently. A real one-step COCO run and its baseline-against-itself control have the same loss bytes, all 369 momentum buffers match bit for bit, and the largest model/buffer difference is 2.98e-8. A 60-step zero-warmup screen diverges in its BF16 trajectory, so this is not a claim of identical stochastic training or final accuracy. The mixed-shape scalar-oracle, storage-budget, alias, DLPack-overlap, dtype, checkpoint-resume and foreach recurrence tests cover the altered scheduling. The full accelerator goldens and long-run accuracy remain unverified; the existing WP-193 status stays open.
The padding region is measured, not universal. A five-repeat optimizer-only crossover screen found the unguarded x scale slower: 111.26 to 119.37 ms per step on the L4. Padding is therefore restricted to a smaller dimension of at most 256, the largest present in the measured nano workload; larger matrices retain exact-shape batches. The guarded screen improved the optimizer in all five scales, but those stage numbers are not the training claim. Parameter subclasses may legitimately refuse storage inspection: a baseline-compatible vector subclass exposed that regression before handoff. Such groups now keep their original keys and scalar momentum, including mixed groups where opaque storage may alias an ordinary parameter.
The final claim includes the complete short fit. Five alternating paired Trainer.fit runs on an isolated NVIDIA L4 (driver 580.178.04, torch 2.9.1+cu129, Python 3.11.17), after two discarded warm-up fits, used detection n, batch 4, 640 px, BF16 and four workers over 256 real COCO val2017 images. Each six-epoch / 384-step fit included cache setup, augmentation, forward/backward, the default three-epoch LR warmup and decay, alpha scheduling, EMA, mosaic closure, four validation batches each epoch and checkpoints. Baseline median 89.448 s (88.765–89.646); candidate 79.183 s (77.799–79.660). The paired reduction median is 11.19% (10.88–13.22%), with the same 809.9 MiB peak allocated. Timing brackets are synchronized and contain no profiler. This is a complete short fit, not a full-COCO or final-accuracy measurement. The author reports that the code measured here passed 65 focused tests on the L4 with two newer-runtime skips, and that those cases passed separately on CUDA. The pytest selection and the revision of that run were not recorded, so the count is not reproducible from this tree; the two optimizer test files collected 71 tests at the reviewed head b060de4. Full make gate and accelerator goldens were not run.
Why the row is lettered 193b. WP-193 had already merged (e8d632c) when this change arrived, and its row stays open on the accelerator-golden check, so the follow-up carries its own row, scope and gate instead of extending the merged one. The commit trailer check accepts digits only, so the commits keep the WP: 193 trailer, as the commits of WP-192b keep WP: 192.