feat(train): add Dion3 as an opt-in matrix optimizer - #1031
Draft
WindChimeRan wants to merge 3 commits into
Draft
Conversation
Contributor
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews
🔴 Require approval from approved reviewers listWaiting for any of
This rule is failing.All pull requests must have at least one approving review from a member of the approved reviewers list before merging.
|
Adds `--optimizer dion3`, substituting Dion3 (microsoft/dion, an alias for
NorDion2) for Muon over the identical parameter split. Muon remains the default;
nothing about the existing paths changes.
Why: Muon's optimizer step is a fixed per-parameter cost that does not shrink
with world size, because Newton-Schulz needs whole matrices and every rank has to
reassemble what FSDP just sharded. Dion3 orthogonalizes only a fraction of the
momentum matrix's rows and megabatches the sharded transfer, so its advantage
grows as you add GPUs. Measured on 8xH100 with DSpark (3 layers, Qwen3-8B
verifier, repo-default 512 anchors), optimizer time per step:
GPUs AdamW Muon Dion3 Muon/Dion3
1 28.8 118.4 37.7 3.14x
2 13.7 164.9 30.3 5.44x
4 8.7 159.7 44.8 3.56x
8 3.1 110.4 10.5 10.52x
End-to-end that is 1.48x at 1 GPU and 1.78x at 8. The win depends on the anchor
budget, since the optimizer cost is constant while forward/backward scale: at
3072 anchors the optimizer is 16-31% of the step rather than 43-58%, and the
end-to-end gain drops to ~1.26x. Peak memory per rank is consistently
Dion3 < Muon < AdamW.
Three implementation notes:
- `dion` is not published on PyPI, so it cannot be a declared dependency (PyPI
rejects direct-URL requirements). It is imported lazily inside the branch with
an actionable error pointing at the git install.
- The device mesh is read off the parameters rather than rebuilt, so it is by
construction the mesh `fully_shard` used. Without it Dion3 silently takes its
single-device path and the multi-GPU benefit disappears.
- `adjust_lr="rms_norm"` is set explicitly. It is 0.2*sqrt(max(fan_out, fan_in)),
the same expression as torch Muon's "match_rms_adamw" default, so a given
--muon-lr means the same effective step size on both. dion's own default
(`spectral_norm`) is a different scale and would silently change the LR.
`dion_selection_scope` defaults to "global" rather than dion's "local": local
takes a per-rank top-k, so the update would depend on how the model is sharded
and results would not reproduce across world sizes. dion's own docs note the two
tie at moderate scale.
Quality is NOT yet established and this is why it ships opt-in: at 300 steps with
untuned hyperparameters, Dion3 reached +1.32% loss and -13.4% acceptance rate
against Muon (AdamW, for reference, is +1.82% loss and -65.6% acceptance).
Hyperparameters need a tuning pass -- dion's CHANGELOG says as much, and the
draft model's 676M orthogonalized parameters sit below the 3B scale at which
dion reports fraction=0.25 to be near-lossless.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Di69UAVwpuJwDSNnEcVpD2
Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
Bisected the inductor miscompile to a torch 2.13.0 regression: 2.12.0 and 2.12.1 both compile the same function cleanly. Comparing the generated Triton for the identical dynamic-shape kernel, 2.12.1's post-loop epilogue re-materializes the gathered value with its own load, while 2.13.0 reuses the temp emitted inside the reduction loop body -- which is a separate scope in Triton, hence `NameError: tmp<N> is not defined`. So the workaround is only needed on >= 2.13. Gating it keeps older supported torch versions (speculators allows >= 2.9.0) on the unmodified path, which matters because pinning shapes static costs extra recompiles -- the exact complaint in microsoft/dion#23. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Di69UAVwpuJwDSNnEcVpD2 Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com>
WindChimeRan
force-pushed
the
feat/dion3-optimizer
branch
from
August 23, 2026 06:15
8058199 to
99a05ff
Compare
This was referenced Aug 23, 2026
|
This pull request has merge conflicts that must be resolved before it can be |
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.
Don't merge
upstream tracking:
Purpose
Aug 12 paper from Tri Dao's group. similar performance, 6x speedup than muon!
https://arxiv.org/pdf/2608.11612
Full claude report: https://claude.ai/code/artifact/c1de5f64-b1e9-4928-8939-280a79521b2c
Qwen3-8B, Training data: 5k ultrachat from
https://huggingface.co/datasets/windchimeran/speculator-tutorial/viewer/default/tutorial_regen
Adds
--optimizer dion3as an opt-in alternative to Muon over the same matrix/scalar parameter split. Muon remains the default, and existing optimizer paths are unchanged.Dion3 is loaded lazily because
microsoft/dionis not published on PyPI.Performance
DSpark, Qwen3-8B verifier, 512 anchors, FSDP2, median of 3 interleaved H100 runs — optimizer time per step:
End-to-end speedup versus Muon was 1.48× on 1 GPU and 1.78× on 8 GPUs. At 3072 anchors, where optimizer time is a smaller portion of each step, the gain dropped to approximately 1.26×.
Quality is not established. An earlier untuned 300-step run—before this PR explicitly matched the scalar AdamW beta defaults—showed 1.32% higher total loss and 13.4% lower acceptance rate than Muon. Those quality results should be rerun; Dion3 therefore remains opt-in and exposes
--dion-fractionfor tuning.Implementation notes
(0.9, 0.999).adjust_lr="rms_norm"to match Muon's effective LR convention.Tests
python -m pytest tests/unit/train/test_dion3_optimizer.py -q: 6 passedpython -m pytest tests/unit/train -q: 303 passed, 5 skippedTested with torch 2.13.0+cu130,
microsoft/dion@58d38ad, and 8×H100 80GB.Checklist
I have filled in: