This repository was archived by the owner on Jul 17, 2026. It is now read-only.
feat(optim): config-gated Muon for the main optimizers (amends SPEC S20)#931
Closed
ocg-goodfire wants to merge 4 commits into
Closed
feat(optim): config-gated Muon for the main optimizers (amends SPEC S20)#931ocg-goodfire wants to merge 4 commits into
ocg-goodfire wants to merge 4 commits into
Conversation
… sharding The prior device_put-onto-sharding was a confirmed no-op (StandardRestore already honors the sharding spec — verified via an instrumented resume: restored CI-fn was correctly ÷N). The real resume-OOM is a ÷1-scale ENTRY RELAYOUT on the first resumed step: orbax-restored arrays carry a default memory layout that differs from what the jitted step was compiled for, so the same executable that runs fresh at ~150GB materialized a ~103GB buffer on resume. The fresh-init reference is built by the same XLA layout assignment as the step, so its `.format` (layout + sharding) IS the step's expected input layout; coercing the restored tree onto it removes the entry relayout. Unblocks durable auto-resume (transient crashes can now recover). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QvFotbQNtDNsgXJQZuzghR
… inputs (S31) (#919) Both hidden-acts eval steps computed leading = residual.shape[:-1] on what is actually the model INPUTS (an LM's [batch, seq] token ids), yielding (batch,) instead of (batch, seq). Every leading-shaped tensor (stochastic delta masks, zero deltas, the clean forward's all-false routes) lost the sequence axis and crashed site_out's route/delta broadcast the first time a run reached a slow eval with the metrics enabled — the deterministic CI step too, via route[..., None] in the clean forward, not just the stochastic delta path. leading now comes off the CI output [*leading, C] (_waist_leading), matching how lm.py::stochastic_site_masks already draws delta masks (ci.shape[:-1]) and staying target-generic. The misnamed residual arg is renamed to inputs: Any per the DecomposedModel protocol — the misnomer is what invited the bug; the sibling attn_patterns_eval.py names the same thing tokens. Regression test with batch != seq (fails pre-fix on both steps) pinning per-site sums/counts and the token-weighted accumulation. Claude-Session: https://claude.ai/code/session_016wiNqVkCMzQ8UQtkeEZ8eK Co-authored-by: Claude Fable 5 <noreply@anthropic.com> (cherry picked from commit ca215e7)
- AnyOptimizerConfig: discriminated adamw|muon union; untyped configs default to adamw (canonical, bit-identical trajectories). - MuonOptimizerConfig: optax.contrib.muon (NS-orthogonalized momentum on 2D leaves, Adam fallback elsewhere), consistent_rms knob, same cosine schedule + S19 clip chain. - Lab canonical assert admits muon (still no-weight-decay subspace). - Tests: muon orthogonalization + adam fallback + discriminator (test_optim_torch_parity), muon checkpoint roundtrip + exact resume (test_checkpoint, S22). - configs/muon/: pile 4L PPGD 40k A/B arms (control + muon 1x/4x lr, consistent_rms 0.2). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This was referenced Jul 11, 2026
Closed
Collaborator
|
Superseded by the combined muon PR #982 (per Oli: one PR for the whole muon feature, rebased onto current feature/jax). All content + validation carried over there. Crew-Address: slack/C08T7UV4449/1783730484.936959 |
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 subscribe to this conversation on GitHub.
Already have an account?
Sign in.
2 participants
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.
What
Adds an experimental, config-gated Muon option for the main (components / ci-fn) optimizers:
AnyOptimizerConfig: discriminatedtype: adamw | muonunion (OptimizerConfigrenamedAdamWOptimizerConfig; a before-validator defaults untyped configs toadamw, so every existing config keeps the bit-identical canonical trajectory).MuonOptimizerConfig→optax.contrib.muon(present in our pinned optax 0.2.8): Newton-Schulz-orthogonalized momentum on 2D leaves (all components leaves are 2D V/U), Adam fallback for non-2D;consistent_rmsknob (0.2 = Kimi scaling, matches update RMS to AdamW's so AdamW LRs transfer). Same cosine schedule, same S19 clip chain.assert_canonical_algorithm_configadmits muon (still constrained to the no-weight-decay subspace).configs/muon/: the pile-4L PPGD 40k A/B arm configs.Includes a cherry-pick of ca215e7 (#919, already on feature/jax — resolves to an empty diff here) because the experiment arms needed it to survive evals.
Why
bridge task try-muon. First A/B on pile 4L PPGD bsc (40k steps, dp=8, seed 0, wandb group
try-muon): Muon Pareto-dominates AdamW — at step 30k, muon@4x-lr reaches CI-L0 257 @ kl_stoch 0.293 vs AdamW's L0 447 @ 0.597 (better on both axes); final train PersistentPGDRecon 14.0 vs 70.2. Muon@matched-RMS-lr already beats AdamW on most metrics, so it isn't purely an LR effect; an adamw@4x-lr disambiguation arm (p-048d4e91) is running — results will be posted here before undrafting.Cost: ~+19% step time on the tiny 4L model (fixed per-site NS overhead; relative cost should shrink at production scale).
Validation
test_optim_torch_parity), muon checkpoint roundtrip + exact resume (test_checkpoint, S22).🤖 Generated with Claude Code