Impose more strict type checking (#86)

* Added torch as additional_dependencies to mypy in pre-commit config.
* Enabled error code explicit-override to mypy.
* Fixed type issues based on AI's recommendation.

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

Author: skyw

.pre-commit-config.yaml

@@ -33,10 +33,11 @@ repos:
     - id: ruff-format
 
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.14.0
+    rev: v1.19.1
     hooks:
       - id: mypy
         exclude: ^docs|^tests|^benchmarks|^docker
+        additional_dependencies: ["torch"]
 
   - repo: local
     hooks:

emerging_optimizers/orthogonalized_optimizers/adaptive_muon.py

@@ -12,7 +12,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Callable, Literal
+from typing import Callable, Literal, overload
 
 
 # TODO(@boxiangw): remove this once bump to python 3.12
@@ -27,6 +27,8 @@
 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
 
 
 class AdaptiveMuon(muon.Muon):
@@ -65,8 +67,8 @@ def __init__(
         *,
         use_nesterov: bool,
         weight_decay_method: opt_mixin.WeightDecayT = "decoupled",
-        fp32_matmul_prec: str,
-        coefficient_type: str = "quintic",
+        fp32_matmul_prec: FP32MatmulPrecT,
+        coefficient_type: NSCoeffT = "quintic",
         num_ns_steps: int = 5,
         scale_mode: muon.MuonScaleT = "spectral",
         extra_scale_factor: float = 1.0,
@@ -179,6 +181,12 @@ def _apply_moment2_normalization(
         else:
             raise TypeError(f"Invalid second moment method: {self.moment2_method}")
 
+    @overload
+    def step(self, closure: None = ...) -> None: ...
+
+    @overload
+    def step(self, closure: Callable[[], float]) -> float: ...
+
     @torch.no_grad()  # type: ignore[misc]
     @override
     def step(self, closure: Callable[[], float] | None = None) -> float | None:

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,13 +50,14 @@ 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:
         def scaled_orthogonalize_fn(grad: torch.Tensor) -> torch.Tensor:
             orth_grad, _, S = polar_via_svd(grad, False)
 
+            scale_factor: float | torch.Tensor
             if scale_mode != "nuclear_norm":
                 scale_factor = muon.get_muon_scale_factor(grad.size(-2), grad.size(-1), mode=scale_mode)
             else:

emerging_optimizers/orthogonalized_optimizers/muon.py

@@ -22,7 +22,9 @@
 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
 
 
 MuonScaleT = Literal["shape_scaling", "spectral", "unit_rms_norm"]
@@ -75,8 +77,8 @@ def __init__(
         *,
         use_nesterov: bool = False,
         weight_decay_method: WeightDecayT = "decoupled",
-        fp32_matmul_prec: str = "medium",
-        coefficient_type: str = "quintic",
+        fp32_matmul_prec: FP32MatmulPrecT = "medium",
+        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

@@ -20,7 +20,9 @@
 from emerging_optimizers import triton_kernels
 
 
-__all__ = ["newton_schulz", "newton_schulz_tp"]
+__all__ = ["newton_schulz", "newton_schulz_tp", "NSCoeffT"]
+
+NSCoeffT = Literal["simple", "quintic", "polar_express", "aol", "custom"]
 
 _COEFFICIENT_SETS = {
     "simple": [
@@ -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,
@@ -121,7 +123,7 @@ def newton_schulz(
     if tp_group is not None:
         X = distributed_normalize_p2(x, eps, tp_group)
     else:
-        X = torch.nn.functional.normalize(x, p=2, dim=(-2, -1), eps=eps)
+        X = torch.nn.functional.normalize(x, p=2, dim=(-2, -1), eps=eps)  # type: ignore[arg-type]
 
     if coefficient_type in _COEFFICIENT_SETS:
         coefficient_sets = _COEFFICIENT_SETS[coefficient_type]
@@ -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/orthogonalized_optimizer.py

@@ -12,7 +12,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Callable
+from typing import Any, Callable, overload
 
 
 # TODO(@boxiangw): remove this once bump to python 3.12
@@ -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,
     ):
@@ -125,6 +126,12 @@ def __init__(
         super().__init__(params, default_args_dict)
         self.scaled_orthogonalize_fn = scaled_orthogonalize_fn
 
+    @overload
+    def step(self, closure: None = ...) -> None: ...
+
+    @overload
+    def step(self, closure: Callable[[], float]) -> float: ...
+
     @torch.no_grad()  # type: ignore[misc]
     @override
     def step(self, closure: Callable[[], float] | None = None) -> float | None:

emerging_optimizers/orthogonalized_optimizers/scion.py

@@ -17,9 +17,11 @@
 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
 
 
 class Scion(OrthogonalizedOptimizer):
@@ -61,8 +63,8 @@ def __init__(
         lr: float = 3e-4,
         momentum_beta: float = 0.95,
         *,
-        fp32_matmul_prec: str = "medium",
-        coefficient_type: str = "quintic",
+        fp32_matmul_prec: FP32MatmulPrecT = "medium",
+        coefficient_type: NSCoeffT = "quintic",
         num_ns_steps: int = 5,
         spectral_radius: float = 1.0,
     ) -> None:
@@ -83,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
 

emerging_optimizers/psgd/procrustes_step.py

@@ -65,15 +65,17 @@ def procrustes_step(
             # rotate Q as exp(a R) Q ~ (I + a R + a^2 R^2/2) Q with an optimal step size by line search
             # for 2nd order expansion, only expand exp(a R) to its 2nd term.
             # Q += _step_size * (RQ + 0.5 * _step_size * RRQ)
-            Q = torch.add(Q, torch.add(RQ, RRQ, alpha=0.5 * step_size), alpha=step_size)
+            Q = torch.add(Q, torch.add(RQ, RRQ, alpha=0.5 * step_size), alpha=step_size)  # type: ignore[call-overload]
         if order == 3:
             RRRQ = R @ RRQ
             tr_RRRQ = torch.trace(RRRQ)
             # for a 3rd order expansion, we take the larger root of the cubic.
             _step_size = (-tr_RRQ - torch.sqrt(tr_RRQ * tr_RRQ - 1.5 * tr_RQ * tr_RRRQ)) / (0.75 * tr_RRRQ)
             step_size = torch.clamp(_step_size, max=max_step_size)
             # Q += step_size * (RQ + 0.5 * step_size * (RRQ + 0.25 * step_size * RRRQ))
-            Q = torch.add(
-                Q, torch.add(RQ, torch.add(RRQ, RRRQ, alpha=0.25 * step_size), alpha=0.5 * step_size), alpha=step_size
+            Q = torch.add(  # type: ignore[call-overload]
+                Q,
+                torch.add(RQ, torch.add(RRQ, RRRQ, alpha=0.25 * step_size), alpha=0.5 * step_size),  # type: ignore[call-overload]
+                alpha=step_size,  # type: ignore[call-overload]
             )
     return Q

emerging_optimizers/psgd/psgd.py

@@ -13,7 +13,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import math
-from typing import Callable, List, Tuple, override
+from typing import Callable, overload
+
+
+try:
+    from typing import override
+except ImportError:
+    from typing_extensions import override
 
 import torch
 from torch.optim.optimizer import ParamsT
@@ -85,6 +91,12 @@ def __init__(
         }
         super().__init__(params, defaults)
 
+    @overload
+    def step(self, closure: None = ...) -> None: ...
+
+    @overload
+    def step(self, closure: Callable[[], float]) -> float: ...
+
     @torch.no_grad()  # type: ignore[misc]
     @override
     def step(self, closure: Callable[[], float] | None = None) -> float | None:
@@ -154,7 +166,7 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
 def _init_psgd_kron_states(
     grad: torch.Tensor,
     precond_init_scale: float = 1.0,
-) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
+) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
     """Initialize the Kronecker factor matrices and Lipschitz constants.
 
     Args:
@@ -165,8 +177,8 @@ def _init_psgd_kron_states(
         q_list: List of Kronecker factors.
         lip_const_list: List of Lipschitz constants for the Kronecker factors.
     """
-    q_list: List[torch.Tensor] = []
-    lip_const_list: List[torch.Tensor] = []
+    q_list: list[torch.Tensor] = []
+    lip_const_list: list[torch.Tensor] = []
 
     # Create identity matrices scaled by precond_init_scale for each dimension
     for size in grad.shape:
@@ -177,13 +189,13 @@ def _init_psgd_kron_states(
 
 
 def _update_precond_procrustes(
-    q_list: List[torch.Tensor],
-    lip_const_list: List[torch.Tensor],
+    q_list: list[torch.Tensor],
+    lip_const_list: list[torch.Tensor],
     exp_avg: torch.Tensor,
     damping_noise_scale: float = 1e-9,
     precond_lr: float = 0.1,
     beta_lip: float = 0.9,
-) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
+) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
     r"""Update the Kron preconditioner Q using procrustes step and uniformization.
 
     Args:
@@ -201,8 +213,8 @@ def _update_precond_procrustes(
     dampened_momentum = exp_avg + (damping_noise_scale + 1e-7 * exp_avg.abs()) * torch.randn_like(exp_avg)
     pg = psgd_kron_contractions.apply_preconditioner(q_list, dampened_momentum)
     total_numel = pg.numel()
-    updated_q_list: List[torch.Tensor] = []
-    updated_lip_const_list: List[torch.Tensor] = []
+    updated_q_list: list[torch.Tensor] = []
+    updated_lip_const_list: list[torch.Tensor] = []
     for dim, q in enumerate(q_list):
         # compute gradient covariance
         precond_grad_cov = psgd_kron_contractions.partial_contraction(pg, pg, dim)
@@ -229,7 +241,7 @@ def _update_matrix_preconditioner(
     total_numel: int,
     precond_lr: float,
     beta_lip: float,
-) -> Tuple[torch.Tensor, torch.Tensor]:
+) -> tuple[torch.Tensor, torch.Tensor]:
     r"""Update matrix-structured preconditioner with adaptive Lipschitz constant.
 
     Args:
@@ -259,7 +271,7 @@ def _update_1d_preconditioner(
     total_numel: int,
     precond_lr: float,
     beta_lip: float,
-) -> Tuple[torch.Tensor, torch.Tensor]:
+) -> tuple[torch.Tensor, torch.Tensor]:
     r"""Update 1D preconditioner with adaptive Lipschitz constant.
 
     Args:

emerging_optimizers/psgd/psgd_utils.py

@@ -12,8 +12,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import List
-
 import torch
 
 
@@ -25,7 +23,7 @@
 
 
 @torch.compile  # type: ignore[misc]
-def uniformize_q_in_place(Q_list: List[torch.Tensor]) -> None:
+def uniformize_q_in_place(Q_list: list[torch.Tensor]) -> None:
     """Balance the dynamic ranges of kronecker factors in place to prevent numerical underflow or overflow.
 
     Each tensor in `Q_list` is rescaled so that its maximum absolute entry