Refactor Muon optimizer constructors

megatron/core/optimizer/emerging_optimizers.py

@@ -23,22 +23,30 @@
 
 try:
     from emerging_optimizers import registry
+    from emerging_optimizers.mixin import WeightDecayT
     from emerging_optimizers.orthogonalized_optimizers import (
         AdaptiveMuon,
+        Moment2MethodT,
         OrthogonalizedOptimizer,
         get_muon_scale_factor,
     )
+    from emerging_optimizers.orthogonalized_optimizers.muon import MuonScaleT
     from emerging_optimizers.orthogonalized_optimizers.muon_utils import NSCoeffT, newton_schulz_tp
 
     # It is necessary to import optimizers for the registry to work.
     from emerging_optimizers.scalar_optimizers import Lion  # pylint: disable=unused-import
     from emerging_optimizers.soap import SOAP  # pylint: disable=unused-import
+    from emerging_optimizers.utils import FP32MatmulPrecT
 
     HAVE_EMERGING_OPTIMIZERS = True
 except ImportError:
     HAVE_EMERGING_OPTIMIZERS = False
     OrthogonalizedOptimizer = object
     AdaptiveMuon = object
+    FP32MatmulPrecT: str
+    Moment2MethodT: str
+    MuonScaleT: str
+    WeightDecayT: str
 
 
 logger = logging.getLogger(__name__)
@@ -157,9 +165,75 @@ def _get_qkv_split_shapes(model_cfg) -> list[int]:
 # ===========================================================================
 
 
+def _tp_muon_pre_init(
+    optim: "TensorParallelMuon",
+    use_decoupled_weight_decay: bool = True,
+    split_qkv: bool = False,
+    is_qkv_fn: Callable[[torch.Tensor], bool] | None = None,
+    qkv_split_shapes: list[int] | None = None,
+    coefficient_type: NSCoeffT = "quintic",
+    num_ns_steps: int = 5,
+    scale_mode: MuonScaleT = "spectral",
+    extra_scale_factor: float = 1.0,
+    pg_collection: Optional[ProcessGroupCollection] = None,
+    tp_mode: Literal["blockwise", "duplicated", "distributed"] = "duplicated",
+) -> tuple[
+    WeightDecayT, Callable[[torch.Tensor, torch.distributed.ProcessGroup, int | None], torch.Tensor]
+]:
+    """Shared initialization functionality of Muon optimizers to support tensor model parallelism.
+
+    Returns the weight decay method and a TP-scaled orthogonalization function.
+    """
+    if num_ns_steps < 1:
+        raise ValueError(f"num_ns_steps must be at least 1, got {num_ns_steps}")
+    if split_qkv and is_qkv_fn is None:
+        raise ValueError("`is_qkv_fn` must not be `None` when `split_qkv=True`")
+
+    def scaled_orthogonalize_fn(
+        grad: torch.Tensor,
+        tp_group: torch.distributed.ProcessGroup,
+        partition_dim: int | None = None,
+    ) -> torch.Tensor:
+        log_single_rank(
+            logger,
+            logging.DEBUG,
+            f'Orthogonalizing grad with {num_ns_steps} steps, '
+            f'{coefficient_type} coefficient, '
+            f'{scale_mode} scale mode, extra_scale_factor={extra_scale_factor}',
+        )
+        size = [grad.size(-2), grad.size(-1)]
+        if partition_dim is not None:
+            size[partition_dim] *= get_pg_size(tp_group)
+        orth_grad = newton_schulz_tp(
+            grad,
+            steps=num_ns_steps,
+            coefficient_type=coefficient_type,
+            tp_group=tp_group,
+            partition_dim=partition_dim,
+            tp_mode="duplicated" if tp_mode == "blockwise" else tp_mode,
+        )
+        scale_factor = get_muon_scale_factor(size[0], size[1], mode=scale_mode)
+        return orth_grad * scale_factor * extra_scale_factor
+
+    optim.pg_collection = pg_collection
+    optim.tp_mode = tp_mode
+    optim.split_qkv = split_qkv
+    optim.is_qkv_fn = is_qkv_fn
+    optim.qkv_split_shapes = qkv_split_shapes
+
+    weight_decay_method = "decoupled" if use_decoupled_weight_decay else "l2"
+    return weight_decay_method, scaled_orthogonalize_fn
+
+
 class TensorParallelMuon(OrthogonalizedOptimizer):
     """Tensor Parallel Muon optimizer."""
 
+    pg_collection: ProcessGroupCollection | None
+    tp_mode: Literal["blockwise", "duplicated", "distributed"]
+    split_qkv: bool
+    is_qkv_fn: Callable[[torch.Tensor], bool] | None
+    qkv_split_shapes: list[int] | None
+
     def __init__(
         self,
         params: ParamsT,
@@ -171,50 +245,29 @@ def __init__(
         split_qkv: bool = False,
         is_qkv_fn: Callable[[torch.Tensor], bool] | None = None,
         qkv_split_shapes: list[int] | None = None,
-        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: str = "spectral",
+        scale_mode: MuonScaleT = "spectral",
         extra_scale_factor: float = 1.0,
         pg_collection: Optional[ProcessGroupCollection] = None,
         tp_mode: Literal["blockwise", "duplicated", "distributed"] = "duplicated",
     ) -> None:
-        if num_ns_steps < 1:
-            raise ValueError(f"num_ns_steps must be at least 1, got {num_ns_steps}")
-
-        def scaled_orthogonalize_fn(
-            grad: torch.Tensor,
-            tp_group: torch.distributed.ProcessGroup,
-            partition_dim: int | None = None,
-        ) -> torch.Tensor:
-            log_single_rank(
-                logger,
-                logging.DEBUG,
-                f'Orthogonalizing grad with {num_ns_steps} steps, '
-                f'{coefficient_type} coefficient, '
-                f'{scale_mode} scale mode, extra_scale_factor={extra_scale_factor}',
-            )
-            size = [grad.size(-2), grad.size(-1)]
-            if partition_dim is not None:
-                size[partition_dim] *= get_pg_size(tp_group)
-            orth_grad = newton_schulz_tp(
-                grad,
-                steps=num_ns_steps,
-                coefficient_type=coefficient_type,
-                tp_group=tp_group,
-                partition_dim=partition_dim,
-                tp_mode="duplicated" if tp_mode == "blockwise" else tp_mode,
-            )
-            scale_factor = get_muon_scale_factor(size[0], size[1], mode=scale_mode)
-            return orth_grad * scale_factor * extra_scale_factor
-
-        self.pg_collection = pg_collection
-        self.tp_mode = tp_mode
-        self.split_qkv = split_qkv
-        self.is_qkv_fn = is_qkv_fn
-        self.qkv_split_shapes = qkv_split_shapes
-
-        weight_decay_method = "decoupled" if use_decoupled_weight_decay else "l2"
+        # When adding common initialization functionality here, please add it to `_tp_muon_pre_init`
+        # so it is also reused by `TensorParallelAdaptiveMuon`.
+        weight_decay_method, scaled_orthogonalize_fn = _tp_muon_pre_init(
+            self,
+            use_decoupled_weight_decay,
+            split_qkv,
+            is_qkv_fn,
+            qkv_split_shapes,
+            coefficient_type,
+            num_ns_steps,
+            scale_mode,
+            extra_scale_factor,
+            pg_collection,
+            tp_mode,
+        )
         # Use explicit class call instead of super() so that subclasses with
         # multiple inheritance (e.g. TensorParallelAdaptiveMuon) don't route
         # through an intermediate class that doesn't accept scaled_orthogonalize_fn.
@@ -369,41 +422,49 @@ def __init__(
         split_qkv: bool = False,
         is_qkv_fn: Callable[[torch.Tensor], bool] | None = None,
         qkv_split_shapes: list[int] | None = None,
-        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: str = "spectral",
+        scale_mode: MuonScaleT = "spectral",
         extra_scale_factor: float = 1.0,
         pg_collection: Optional[ProcessGroupCollection] = None,
         tp_mode: Literal["blockwise", "duplicated", "distributed"] = "duplicated",
-        moment2_method: Literal["adamuon", "normuon"] = "adamuon",
+        moment2_method: Moment2MethodT = "adamuon",
         beta2: float = 0.95,
         eps: float = 1e-8,
     ) -> None:
-        TensorParallelMuon.__init__(
+        weight_decay_method, scaled_orthogonalize_fn = _tp_muon_pre_init(
+            self,
+            use_decoupled_weight_decay,
+            split_qkv,
+            is_qkv_fn,
+            qkv_split_shapes,
+            coefficient_type,
+            num_ns_steps,
+            scale_mode,
+            extra_scale_factor,
+            pg_collection,
+            tp_mode,
+        )
+        AdaptiveMuon.__init__(
             self,
             params,
             lr=lr,
             momentum=momentum,
-            nesterov=nesterov,
             weight_decay=weight_decay,
-            use_decoupled_weight_decay=use_decoupled_weight_decay,
-            split_qkv=split_qkv,
-            is_qkv_fn=is_qkv_fn,
-            qkv_split_shapes=qkv_split_shapes,
+            nesterov=nesterov,
+            weight_decay_method=weight_decay_method,
             fp32_matmul_prec=fp32_matmul_prec,
             coefficient_type=coefficient_type,
             num_ns_steps=num_ns_steps,
             scale_mode=scale_mode,
             extra_scale_factor=extra_scale_factor,
-            pg_collection=pg_collection,
-            tp_mode=tp_mode,
+            moment2_method=moment2_method,
+            beta2=beta2,
+            eps=eps,
         )
-        self.moment2_method = moment2_method
-
-        for group in self.param_groups:
-            group.setdefault("beta2", beta2)
-            group.setdefault("eps", eps)
+        # AdaptiveMuon sets its own `scaled_orthogonalize_fn`, so we override it afterwards.
+        self.scaled_orthogonalize_fn = scaled_orthogonalize_fn
 
     @torch.no_grad()  # type: ignore[misc]
     def step(self, closure: Optional[Callable] = None) -> Optional[float]:

megatron/core/optimizer/optimizer_config.py

@@ -300,9 +300,6 @@ class OptimizerConfig:
     soap_shampoo_beta: float = 0.95
     """The beta parameter for the Shampoo preconditioner."""
 
-    soap_precondition_frequency: int = 1
-    """The frequency of the Shampoo preconditioner."""
-
     soap_use_kl_shampoo: bool = True
     """Whether to use the KL-Shampoo preconditioner."""
 

pyproject.toml

@@ -231,7 +231,7 @@ flash_mla = [
 ]
 transformer-engine = { git = "https://github.com/NVIDIA/TransformerEngine.git", rev = "e7c550c5f80636cf841a8204b1d6f85a5f3f28b7" }
 nemo-run = { git = "https://github.com/NVIDIA-NeMo/Run.git", rev = "ddd40a8f24847f5c919f911d0240bd622653612f" }
-emerging_optimizers = { git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git", rev = "v0.2.0" }
+emerging_optimizers = { git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git", rev = "v0.3.0" }
 fast-hadamard-transform = { git = "https://github.com/Dao-AILab/fast-hadamard-transform.git", rev = "f134af63deb2df17e1171a9ec1ea4a7d8604d5ca" }
 mamba-ssm = { git = "https://github.com/state-spaces/mamba.git", rev = "0048fbf2e7b2f214dcbe703ea3dec2b9647595e1" }
 

tests/unit_tests/test_emerging_optimizers.py

@@ -1320,12 +1320,7 @@ def test_soap_optimizer_smoke():
     model.weight.data.fill_(1.0)
 
     optimizer = SOAP(
-        params=[model.weight],
-        lr=0.01,
-        betas=(0.9, 0.999),
-        shampoo_beta=0.95,
-        weight_decay=0.01,
-        precondition_frequency=1,
+        params=[model.weight], lr=0.01, betas=(0.9, 0.999), shampoo_beta=0.95, weight_decay=0.01
     )
 
     # Test basic properties
@@ -1373,12 +1368,7 @@ def test_soap_optimizer_multiple_steps():
     model.weight.data.fill_(1.0)
 
     optimizer = SOAP(
-        params=[model.weight],
-        lr=0.01,
-        betas=(0.9, 0.999),
-        shampoo_beta=0.95,
-        weight_decay=0.01,
-        precondition_frequency=1,
+        params=[model.weight], lr=0.01, betas=(0.9, 0.999), shampoo_beta=0.95, weight_decay=0.01
     )
 
     weights_history = [model.weight.data.clone()]
@@ -1400,36 +1390,6 @@ def test_soap_optimizer_multiple_steps():
         ), f"Weight should change at step {i}"
 
 
-@skip_no_soap
-@pytest.mark.parametrize("precondition_frequency", [1, 5, 10])
-def test_soap_optimizer_precondition_frequency(precondition_frequency):
-    """Test SOAP optimizer with different precondition frequencies."""
-
-    model = torch.nn.Linear(60, 30, bias=False, dtype=torch.float32, device='cuda')
-    model.requires_grad_(True)
-    model.weight.data.fill_(1.0)
-
-    optimizer = SOAP(
-        params=[model.weight],
-        lr=0.01,
-        betas=(0.9, 0.999),
-        shampoo_beta=0.95,
-        precondition_frequency=precondition_frequency,
-    )
-
-    input_tensor = torch.randn(16, 60, dtype=torch.float32, device='cuda')
-    output = model(input_tensor)
-    loss = output.sum()
-    loss.backward()
-
-    original_weight = model.weight.data.clone()
-    optimizer.step()
-
-    assert not torch.equal(
-        model.weight.data, original_weight
-    ), f"Weight should be updated with precondition_frequency={precondition_frequency}"
-
-
 @skip_no_soap
 @pytest.mark.parametrize("use_kl_shampoo", [True, False])
 def test_soap_optimizer_kl_shampoo(use_kl_shampoo):
@@ -1445,7 +1405,6 @@ def test_soap_optimizer_kl_shampoo(use_kl_shampoo):
         betas=(0.9, 0.999),
         shampoo_beta=0.95,
         use_kl_shampoo=use_kl_shampoo,
-        precondition_frequency=1,
     )
 
     input_tensor = torch.randn(16, 60, dtype=torch.float32, device='cuda')
@@ -1470,13 +1429,7 @@ def test_soap_optimizer_shampoo_beta(shampoo_beta):
     model.requires_grad_(True)
     model.weight.data.fill_(1.0)
 
-    optimizer = SOAP(
-        params=[model.weight],
-        lr=0.01,
-        betas=(0.9, 0.999),
-        shampoo_beta=shampoo_beta,
-        precondition_frequency=1,
-    )
+    optimizer = SOAP(params=[model.weight], lr=0.01, betas=(0.9, 0.999), shampoo_beta=shampoo_beta)
 
     input_tensor = torch.randn(16, 60, dtype=torch.float32, device='cuda')
     output = model(input_tensor)
@@ -1527,7 +1480,6 @@ def test_get_megatron_optimizer_soap_smoke(self):
             bf16=True,
             use_distributed_optimizer=False,
             soap_shampoo_beta=0.95,
-            soap_precondition_frequency=1,
             soap_use_kl_shampoo=True,
         )