Repository navigation
feat: segmented CUDA fast-clip + sharded-AdamW memory fix + docs - #13
Merged
Merged
Conversation
Add an optional CUDA fast path for segment-local gradient clipping (regular / muon / adamw buckets), with a distributed-safe design: - Three stages -- local per-bucket squared norm, one rank-uniform all_reduce, local in-place scaling. CUDA kernels accelerate eligible tensors; ineligible tensors (non-contiguous, unusual dtype, CPU) or a missing extension fall back to torch per tensor. The fast/slow choice is purely local, so it never changes collective count/shape and stays safe to call from every rank in lockstep. - The clip coefficient is computed in torch (sqrt + clamp) rather than a CUDA kernel, so NaN/inf handling matches the reference exactly. - torch is intentionally not a build dependency: pinning it would force isolated builds to download a generic torch and risk an ABI-mismatched .so. Compile with `--no-build-isolation`; otherwise the pure-Python clip is used. - The extension loader distinguishes "not built" (quiet) from "built but failed to load" (warns once), so silent slow-path fallback is diagnosable. DMUON_FAST_CLIP_VERBOSE=1 re-raises. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
… guard - prepare_muon_grads/_ensure_grads_ready gain a coalesce_wait path: when every group shares one reduce stream, issue a single aggregate stream wait instead of draining each group individually. Used by the fused fast-clip path; defaults preserve the per-group wait. - sharded_adamw params no longer allocate a redundant full fp32 owner copy (_owned_data): every shard rank already owns its flat shard in _sharded_adamw_data and reconstructs the forward view via all-gather. - Guard the param-group route override: moving a param constructed on the sharded_adamw route to another route would leave later steps dereferencing a None _owned_data (crash) or silently skipping the update. Reject the move with a clear error instead. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
- Unit tests: chunk-size resolution, CPU segment-local semantics, non-contiguous tensors clip (not skipped), CUDA-vs-torch parity including NaN/inf, and optimizer stat wiring. - Regression for the distributed-hang: a mocked-dist test asserts the clip issues exactly one all_reduce of shape [num_buckets] regardless of a rank-local non-finite value. - A 2-rank gloo torchrun script reproduces "inf on one rank only" and asserts the clip completes (no hang) with correct all-reduced norms. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
- Installation: how the optional CUDA extension is compiled (--no-build-isolation), the env toggles, and the load-failure warning. - API reference: the segmented bucket-clip public API. - Training guide: pointer to the fused regular/muon/adamw clip. English and Chinese kept in parity. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
- Point site_url/repo_url and all in-doc links at the X-Square-Robot org (the GitHub Pages site is served from the org, not the personal account the config previously named -- fixing canonical/sitemap links). - README: real informative badges incl. a docs badge to the live site, plus a documentation quick-links nav block. Co-Authored-By: Claude Opus 4.8 <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.
Summary
Adds an optional segmented CUDA fast path for gradient clipping (regular / muon / adamw buckets), a sharded-AdamW memory fix, a coalesced grad-wait optimization, and publishes the full docs site.
The fast-clip path is designed distributed-safe from the ground up: local per-bucket norm → one rank-uniform all_reduce → local scaling. CUDA kernels accelerate eligible tensors; ineligible tensors (or a missing extension) fall back to torch per tensor, so the fast/slow choice never changes collective count or
shape and stays safe to call from every rank in lockstep.
Commits
Highlights
Testing