add a liberal for matmul precision

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.utils import FP32MatmulPrecT
 
 
 class AdaptiveMuon(muon.Muon):
@@ -65,7 +66,7 @@ def __init__(
         *,
         use_nesterov: bool,
         weight_decay_method: opt_mixin.WeightDecayT = "decoupled",
-        fp32_matmul_prec: str,
+        fp32_matmul_prec: FP32MatmulPrecT,
         coefficient_type: str = "quintic",
         num_ns_steps: int = 5,
         scale_mode: muon.MuonScaleT = "spectral",

emerging_optimizers/orthogonalized_optimizers/mop.py

@@ -22,6 +22,7 @@
 from emerging_optimizers.mixin import WeightDecayT
 from emerging_optimizers.orthogonalized_optimizers import muon
 from emerging_optimizers.orthogonalized_optimizers.orthogonalized_optimizer import OrthogonalizedOptimizer, _args_doc
+from emerging_optimizers.utils import FP32MatmulPrecT
 
 
 __all__ = ["MOP"]
@@ -49,7 +50,7 @@ def __init__(
         *,
         use_nesterov: bool = False,
         weight_decay_method: WeightDecayT = "decoupled",
-        fp32_matmul_prec: str = "highest",
+        fp32_matmul_prec: FP32MatmulPrecT = "highest",
         scale_mode: muon.MuonScaleT | Literal["nuclear_norm"] = "nuclear_norm",
         extra_scale_factor: float = 1.0,
     ) -> None:

emerging_optimizers/orthogonalized_optimizers/muon.py

@@ -23,6 +23,7 @@
 from emerging_optimizers.mixin import WeightDecayT
 from emerging_optimizers.orthogonalized_optimizers import muon_utils
 from emerging_optimizers.orthogonalized_optimizers.orthogonalized_optimizer import OrthogonalizedOptimizer, _args_doc
+from emerging_optimizers.utils import FP32MatmulPrecT
 
 
 MuonScaleT = Literal["shape_scaling", "spectral", "unit_rms_norm"]
@@ -75,7 +76,7 @@ def __init__(
         *,
         use_nesterov: bool = False,
         weight_decay_method: WeightDecayT = "decoupled",
-        fp32_matmul_prec: str = "medium",
+        fp32_matmul_prec: FP32MatmulPrecT = "medium",
         coefficient_type: str = "quintic",
         num_ns_steps: int = 5,
         scale_mode: MuonScaleT = "spectral",

emerging_optimizers/orthogonalized_optimizers/orthogonalized_optimizer.py

@@ -28,6 +28,7 @@
 
 from emerging_optimizers import mixin as opt_mixin
 from emerging_optimizers import utils
+from emerging_optimizers.utils import FP32MatmulPrecT
 
 
 _args_doc = """params: Iterable of parameters to optimize or dicts defining parameter groups
@@ -103,7 +104,7 @@ def __init__(
         *,
         use_nesterov: bool,
         weight_decay_method: opt_mixin.WeightDecayT,
-        fp32_matmul_prec: str,
+        fp32_matmul_prec: FP32MatmulPrecT,
         scaled_orthogonalize_fn: Callable | None = None,
         **kwargs: Any,
     ):

emerging_optimizers/orthogonalized_optimizers/scion.py

@@ -20,6 +20,7 @@
 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.orthogonalized_optimizer import OrthogonalizedOptimizer
+from emerging_optimizers.utils import FP32MatmulPrecT
 
 
 class Scion(OrthogonalizedOptimizer):
@@ -61,7 +62,7 @@ def __init__(
         lr: float = 3e-4,
         momentum_beta: float = 0.95,
         *,
-        fp32_matmul_prec: str = "medium",
+        fp32_matmul_prec: FP32MatmulPrecT = "medium",
         coefficient_type: str = "quintic",
         num_ns_steps: int = 5,
         spectral_radius: float = 1.0,

emerging_optimizers/soap/soap.py

@@ -31,6 +31,7 @@
 from emerging_optimizers import mixin as opt_mixin
 from emerging_optimizers import scalar_optimizers, utils
 from emerging_optimizers.soap import soap_utils
+from emerging_optimizers.utils import FP32MatmulPrecT
 
 
 __all__ = [
@@ -97,9 +98,9 @@ def __init__(
         adam_warmup_steps: int = 0,
         precondition_1d: bool = False,
         correct_bias: bool = True,
-        fp32_matmul_prec: str = "high",
+        fp32_matmul_prec: FP32MatmulPrecT = "high",
         use_eigh: bool = False,
-        qr_fp32_matmul_prec: str = "high",
+        qr_fp32_matmul_prec: FP32MatmulPrecT = "high",
         use_adaptive_criteria: bool = False,
         adaptive_update_tolerance: float = 1e-7,
         power_iter_steps: int = 1,

emerging_optimizers/utils/__init__.py

@@ -13,18 +13,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from contextlib import contextmanager
-from typing import Generator
+from typing import Generator, Literal
 
 import torch
 
 from .eig import *
 
 
-__all__ = ["fp32_matmul_precision", "get_pg_size", "get_pg_rank"]
+__all__ = ["fp32_matmul_precision", "get_pg_size", "get_pg_rank", "FP32MatmulPrecT"]
+
+FP32MatmulPrecT = Literal["highest", "high", "medium"]
 
 
 @contextmanager
-def fp32_matmul_precision(precision: str = "highest") -> Generator[None, None, None]:
+def fp32_matmul_precision(precision: FP32MatmulPrecT = "highest") -> Generator[None, None, None]:
     """Context manager for setting the precision of matmuls.
 
     Args: