Repository navigation
Sync-free clip stats, foreach AdamW, NS CUDA-graph replay, SYRK wiring fix - #16
Merged
Merged
Conversation
MuonGradClipStats and GradClipBucketStats now always carry total_norm / clip_coef / clipped / found_inf as 0-dim device tensors: the clip call performs no device-to-host sync (the former per-step packed copy in _build_stats and the eager .item() reads are gone). Consumers that need python scalars (logging) materialize on their own schedule with one batched readback; the only reason to read these on the host at clip time would be host-side control flow (inf-skip / dynamic thresholds), which DMuon does not do. The DMUON_CLIP_DEFER_HOST_STATS switch is removed -- deferred is the only path. Also fixes the clipped semantics: finiteness is judged on the NORM, not the coefficient. An inf norm drives the coef to 0.0 (finite!), so the old isfinite(coef) check reported clipped=True on inf gradients; now inf reports found_inf=True, clipped=False. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Both AdamW paths (dedicated owner-managed and FSDP2 symmetric) now collect (state_key, param, grad) entries and update through one foreach-batched implementation; the per-tensor _adamw_update_tensor is removed. A per-param loop launched ~8 kernels per tensor (700+ launches/step for ~90 owned adamw params) with python state lookups between; each foreach bucket is a dozen multi-tensor launches. Entries bucket by (post-increment step, param dtype/device, grad dtype): bias corrections are per-STEP scalars and params that historically skipped steps (grad=None) carry diverged counters, so bucketing keeps their correction exact. Op order matches the removed per-tensor path exactly (decoupled wd -> moments -> bias-corrected addcdiv). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
One NS invocation issues ~25-30 kernels through python; a rank stepping ~50 muon params drives >1000 launches per step -- the bulk of the CPU- serial optimizer section. NS is branch-free with a small stable set of input shapes (11-20 per rank), so capture it once per (shape, dtype, device) and replay: measured host cost per NS call drops 2.5ms -> 26us, replay is bitwise-identical to eager. Design points: a persistent stage_in buffer decouples the graph from the step-fresh grad allocations; lr/wd/momentum stay in the eager pre/post steps so schedule changes never re-capture; the first call per shape runs eagerly (it doubles as SYRK-autotune/JIT warmup, which must not be captured), the second captures; TP params never enter graphs (collective- managed buffer lifetimes); capture failure degrades loudly to eager for that shape only. DMUON_NS_CUDA_GRAPH=0 opts out. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
NewtonSchulz resolved kernel=quack (and advertised it via get_ns_backend) but syrk_or_cublas never consulted the choice -- it always autotuned the SM80 cute kernel against cuBLAS, so on SM90 the quack path silently never ran. NewtonSchulz.__init__ now publishes the resolved kernel to the dispatch layer, and syrk_or_cublas routes: quack goes through a per-shape quack-vs-cuBLAS autotune (numerical cross-check against cuBLAS first, loud reject beyond tolerance), kernel=cublas keeps its bit-exact direct path, anything else keeps the legacy cute autotune. Autotune timing switches to a pre-queued GPU-execution-only measurement: a blocker kernel (torch.cuda._sleep, guarded for the private API) keeps the GPU busy while the host enqueues all iterations, so the timed window contains only back-to-back kernel execution. Rationale: the NS hot path is moving to CUDA-graph replay where per-call host cost is ~zero, so backend selection must compare kernels, not python wrappers (the quack high-level interface costs ~40us host per call, which used to flip the 2560x2560 verdict to cuBLAS even though the quack kernel is 1.4x faster; L2 hot/cold can additionally swing cuBLAS timings 3x). Persisted autotune cache filenames gain a _v2 suffix so wall-clock-era verdicts are not reused after upgrade. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- guides/training.md: clip stats are 0-dim device tensors (no D2H at clip time; materialize on your own logging schedule; inf reports found_inf=True, clipped=False), bucket stats follow the same contract. - reference/newton-schulz.md: quack now genuinely dispatches (wiring fix) with the per-shape numerics-gated GPU-only autotune; new CUDA-graph replay section (default on, opt-out, shape-stability assumption, first-step capture cost, TP boundary, failure degradation). - troubleshooting.md: _v2 autotune-cache generation note (one-time re-autotune after upgrade). - zh/ mirrors updated to parity. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Four independent host-overhead eliminations for the optimizer hot path, plus a
correctness fix in kernel dispatch.
Commits
perf(clip): device-tensor stats, no D2H in the clip path — stats are
0-dim CUDA tensors; consumers materialize on their logging schedule. Also
fixes
clippedon inf gradients (finiteness now judged on the norm; inf →found_inf=True, clipped=False).perf(optim): foreach-only AdamW — one batched implementation for both
dedicated and FSDP2-symmetric params; ~700+ launches/step collapse to a
dozen per bucket; per-step bucketing keeps bias correction exact for params
with diverged step counters.
perf(optim): per-shape CUDA-graph replay for Newton-Schulz (default on) —
host cost per NS call 2.5ms → 26µs, replay bitwise-identical to eager;
first call per shape stays eager (autotune/JIT warmup), TP params excluded,
loud per-shape fallback on capture failure.
DMUON_NS_CUDA_GRAPH=0opts out.fix(optim): wire the resolved SYRK kernel; GPU-only autotune —
kernel="quack"previously never reachedsyrk_or_cublas(always fell tothe cute-vs-cuBLAS autotune). Now wired, with a per-shape numerics-gated
quack-vs-cuBLAS autotune measured in GPU-execution-only time (pre-queued
behind a blocker kernel): with graph replay the per-call host cost is ~zero,
so selection must compare kernels, not python wrappers. Cache files gain a
_v2generation suffix.docs: guides/reference/troubleshooting (+zh mirrors) updated to match.
🤖 Generated with Claude Code