add a liberal for NS coefficient

Signed-off-by: Hao Wu <[email protected]>

Author: skyw

emerging_optimizers/orthogonalized_optimizers/adaptive_muon.py

@@ -27,6 +27,7 @@
 from emerging_optimizers import mixin as opt_mixin
 from emerging_optimizers import utils
 from emerging_optimizers.orthogonalized_optimizers import muon
+from emerging_optimizers.orthogonalized_optimizers.muon_utils import NSCoeffT
 from emerging_optimizers.utils import FP32MatmulPrecT
 
 
@@ -67,7 +68,7 @@ def __init__(
         use_nesterov: bool,
         weight_decay_method: opt_mixin.WeightDecayT = "decoupled",
         fp32_matmul_prec: FP32MatmulPrecT,
-        coefficient_type: str = "quintic",
+        coefficient_type: NSCoeffT = "quintic",
         num_ns_steps: int = 5,
         scale_mode: muon.MuonScaleT = "spectral",
         extra_scale_factor: float = 1.0,

emerging_optimizers/orthogonalized_optimizers/muon.py

@@ -22,6 +22,7 @@
 from emerging_optimizers import triton_kernels
 from emerging_optimizers.mixin import WeightDecayT
 from emerging_optimizers.orthogonalized_optimizers import muon_utils
+from emerging_optimizers.orthogonalized_optimizers.muon_utils import NSCoeffT
 from emerging_optimizers.orthogonalized_optimizers.orthogonalized_optimizer import OrthogonalizedOptimizer, _args_doc
 from emerging_optimizers.utils import FP32MatmulPrecT
 
@@ -77,7 +78,7 @@ def __init__(
         use_nesterov: bool = False,
         weight_decay_method: WeightDecayT = "decoupled",
         fp32_matmul_prec: FP32MatmulPrecT = "medium",
-        coefficient_type: str = "quintic",
+        coefficient_type: NSCoeffT = "quintic",
         num_ns_steps: int = 5,
         scale_mode: MuonScaleT = "spectral",
         extra_scale_factor: float = 1.0,

emerging_optimizers/orthogonalized_optimizers/muon_utils.py

@@ -22,6 +22,8 @@
 
 __all__ = ["newton_schulz", "newton_schulz_tp"]
 
+NSCoeffT = Literal["simple", "quintic", "polar_express", "aol", "custom"]
+
 _COEFFICIENT_SETS = {
     "simple": [
         (3.4445, -4.7750, 2.0315),
@@ -67,7 +69,7 @@ def distributed_normalize_p2(x: torch.Tensor, eps: float, group: torch.distribut
 def newton_schulz(
     x: torch.Tensor,
     steps: int,
-    coefficient_type: str = "quintic",
+    coefficient_type: NSCoeffT = "quintic",
     custom_coefficient_sets: list[tuple[float, float, float]] | None = None,
     eps: float = 1e-7,
     transpose: bool | None = None,
@@ -164,7 +166,7 @@ def newton_schulz(
 def newton_schulz_tp(
     x: torch.Tensor,
     steps: int,
-    coefficient_type: str,
+    coefficient_type: NSCoeffT,
     tp_group: torch.distributed.ProcessGroup,
     partition_dim: int | None = None,
     mode: Literal["duplicated", "distributed"] = "duplicated",

emerging_optimizers/orthogonalized_optimizers/scion.py

@@ -17,8 +17,9 @@
 from absl import logging
 from torch.optim.optimizer import ParamsT
 
+from emerging_optimizers.orthogonalized_optimizers import muon_utils
 from emerging_optimizers.orthogonalized_optimizers.muon import get_muon_scale_factor
-from emerging_optimizers.orthogonalized_optimizers.muon_utils import newton_schulz
+from emerging_optimizers.orthogonalized_optimizers.muon_utils import NSCoeffT
 from emerging_optimizers.orthogonalized_optimizers.orthogonalized_optimizer import OrthogonalizedOptimizer
 from emerging_optimizers.utils import FP32MatmulPrecT
 
@@ -63,7 +64,7 @@ def __init__(
         momentum_beta: float = 0.95,
         *,
         fp32_matmul_prec: FP32MatmulPrecT = "medium",
-        coefficient_type: str = "quintic",
+        coefficient_type: NSCoeffT = "quintic",
         num_ns_steps: int = 5,
         spectral_radius: float = 1.0,
     ) -> None:
@@ -84,7 +85,9 @@ def scaled_orthogonalize_fn(grad: torch.Tensor) -> torch.Tensor:
             logging.debug(
                 f"Orthogonalizing grad with {num_ns_steps} steps, {coefficient_type} coefficient, spectral_radius={spectral_radius}"
             )
-            orth_grad = newton_schulz(grad, steps=num_ns_steps, coefficient_type=coefficient_type, use_syrk=False)
+            orth_grad = muon_utils.newton_schulz(
+                grad, steps=num_ns_steps, coefficient_type=coefficient_type, use_syrk=False
+            )
             width_factor = get_muon_scale_factor(grad.size(-2), grad.size(-1), mode="unit_rms_norm")
             return orth_grad * width_factor * spectral_radius