bump emerging optimizers to v0.3.0

megatron/core/optimizer/emerging_optimizers.py

@@ -399,6 +399,8 @@ def __init__(
             pg_collection=pg_collection,
             tp_mode=tp_mode,
         )
+        self.scale_mode = scale_mode
+        self.extra_scale_factor = extra_scale_factor
         self.moment2_method = moment2_method
 
         for group in self.param_groups:

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,
         )