Make Nora+NAdamW the default optimizer

2026-07-31 23:16:39 +08:00
parent 7aa5ed09d9
commit 04899a2b15
11 changed files with 1010 additions and 105 deletions
@@ -0,0 +1,36 @@
+"""Optimizer implementations and factory registration."""
+
+from torch.optim import Optimizer
+
+from astrai.factory import BaseFactory
+
+
+class OptimizerFactory(BaseFactory[Optimizer]):
+    """Factory for built-in training optimizers."""
+
+
+from astrai.optim.muon_mix import MuonMix
+from astrai.optim.nora_nadamw import (
+    NAdamW,
+    Nora,
+    NoraNAdamW,
+    OptimizerParameterGroups,
+    nora_direction,
+    nora_lr_scale,
+    partition_optimizer_parameters,
+)
+
+OptimizerFactory.register("nora_nadamw")(NoraNAdamW)
+OptimizerFactory.register("muon_adamw")(MuonMix)
+
+__all__ = [
+    "MuonMix",
+    "NAdamW",
+    "Nora",
+    "NoraNAdamW",
+    "OptimizerFactory",
+    "OptimizerParameterGroups",
+    "nora_direction",
+    "nora_lr_scale",
+    "partition_optimizer_parameters",
+]