ViperEkura
6db276f37a
- implement Mano (v2) with axis-rotating tangent projection and manifold normalization, replacing Newton-Schulz iteration - composite ManoAdamW reuses partition_optimizer_parameters and composite helpers - register mano_adamw in OptimizerFactory, export Mano and ManoAdamW - add --mano_momentum and --mano_nesterov CLI options in Optimizer group - add mano_adamw hyperparameters branch in train.py - document mano_adamw in params.md - add tests for single-step projection, axis alternation, factory registration, closure, and resume
39 lines
846 B
Python
39 lines
846 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.mano_adamw import Mano, ManoAdamW
|
|
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__ = [
|
|
"Mano",
|
|
"ManoAdamW",
|
|
"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",
|
|
]
|