ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: extract composite optimizer helpers and unify naming

Browse Source 
- add astrai/optim/composite.py with shared step/zero_grad/state_dict/param_groups helpers and OptimizerFactory
- rename MuonMix to MuonAdamW (matches registered name muon_adamw) and file to muon_adamw.py
- use @OptimizerFactory.register decorator in each optimizer module instead of post-import registration in __init__
- fix closure being invoked once per sub-optimizer in MuonAdamW.step (now exactly once via composite_step)
- NoraNAdamW.step now forwards closure correctly
This commit is contained in:
ViperEkura
2026-08-01 08:07:45 +08:00
parent 25c9e81b2b
commit 11073bd1d2
5 changed files with 119 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/guides/params.md
+1 -1
View File
@@ -56,7 +56,7 @@ under DTensor sharding and rejects layouts sharded along the last dimension.
| `--nora_weight_decay` | Nora matrix weight decay | 0.0 |
Optimizer identity and hyperparameters are saved in checkpoint metadata. Optimizer
states are intentionally not interchangeable: resume older MuonMix checkpoints
states are intentionally not interchangeable: resume older MuonAdamW checkpoints
with `--optimizer=muon_adamw`.
### Data Loading