ViperEkura
11073bd1d2
- 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
36 lines
765 B
Python
36 lines
765 B
Python
"""Optimizer implementations and factory registration."""
|
|
|
|
from astrai.optim.composite import (
|
|
OptimizerFactory,
|
|
composite_state_dict,
|
|
composite_step,
|
|
composite_zero_grad,
|
|
refresh_param_groups,
|
|
)
|
|
from astrai.optim.muon_adamw import MuonAdamW
|
|
from astrai.optim.nora_nadamw import (
|
|
NAdamW,
|
|
Nora,
|
|
NoraNAdamW,
|
|
OptimizerParameterGroups,
|
|
nora_direction,
|
|
nora_lr_scale,
|
|
partition_optimizer_parameters,
|
|
)
|
|
|
|
__all__ = [
|
|
"MuonAdamW",
|
|
"NAdamW",
|
|
"Nora",
|
|
"NoraNAdamW",
|
|
"OptimizerFactory",
|
|
"OptimizerParameterGroups",
|
|
"composite_state_dict",
|
|
"composite_step",
|
|
"composite_zero_grad",
|
|
"nora_direction",
|
|
"nora_lr_scale",
|
|
"partition_optimizer_parameters",
|
|
"refresh_param_groups",
|
|
]
|