update type opt

PrimeIntellect-ai · Dec 6, 2024 · bcbf8fd · bcbf8fd
1 parent 30a0f23
commit bcbf8fd
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/src/zeroband/optimizers/__init__.py b/src/zeroband/optimizers/__init__.py
@@ -1,8 +1,12 @@
+from typing import TypeAlias
 import torch
 from zeroband.optimizers.muon import Muon, AdamConfig, MuonConfig
 
 
-def get_optimizer(params: list[torch.nn.Parameter], config: AdamConfig | MuonConfig) -> torch.optim.Optimizer:
+OptimizersConfig: TypeAlias = AdamConfig | MuonConfig
+
+
+def get_optimizer(params: list[torch.nn.Parameter], config: OptimizersConfig) -> torch.optim.Optimizer:
     if isinstance(config, AdamConfig):
         return torch.optim.AdamW(
             params,
@@ -25,4 +29,4 @@ def get_optimizer(params: list[torch.nn.Parameter], config: AdamConfig | MuonCon
         raise ValueError(f"Unknown optimizer {config.optimizer}")
 
 
-__all__ = ["AdamConfig", "MuonConfig", "get_optimizer"]
+__all__ = ["OptimizersConfig", "get_optimizer"]
diff --git a/src/zeroband/train.py b/src/zeroband/train.py
@@ -19,7 +19,7 @@
 from zeroband.diloco import Diloco, DilocoConfig
 from zeroband.comms import ElasticDeviceMesh
 from zeroband.loss import cross_entropy_max_z_loss
-from zeroband.optimizers import AdamConfig, MuonConfig, get_optimizer
+from zeroband.optimizers import AdamConfig, OptimizersConfig, get_optimizer
 
 from zeroband.utils import (
     FakeTokenizer,
@@ -42,7 +42,7 @@
 
 
 class OptimConfig(BaseConfig):
-    optim: AdamConfig | MuonConfig = AdamConfig()
+    optim: OptimizersConfig = AdamConfig()
 
     sched_type: Literal["cosine", "linear", "wsd-sqrt"] = "cosine"
     warmup_steps: int = 1000