diff --git a/megatron/core/optimizer/emerging_optimizers.py b/megatron/core/optimizer/emerging_optimizers.py index 53ac956b35c..7e936f719eb 100644 --- a/megatron/core/optimizer/emerging_optimizers.py +++ b/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: diff --git a/megatron/core/optimizer/optimizer_config.py b/megatron/core/optimizer/optimizer_config.py index 24f9a032c47..f19e9970192 100644 --- a/megatron/core/optimizer/optimizer_config.py +++ b/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.""" diff --git a/pyproject.toml b/pyproject.toml index e19fe870334..764aec8d3a1 100644 --- a/pyproject.toml +++ b/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" } diff --git a/tests/unit_tests/test_emerging_optimizers.py b/tests/unit_tests/test_emerging_optimizers.py index e3b9f666fb2..7f0aed21329 100644 --- a/tests/unit_tests/test_emerging_optimizers.py +++ b/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, ) diff --git a/uv.lock b/uv.lock index 15036ceb4ce..cc32e52243e 100644 --- a/uv.lock +++ b/uv.lock @@ -1115,8 +1115,8 @@ wheels = [ [[package]] name = "emerging-optimizers" -version = "0.2.0" -source = { git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.2.0#1effa026ff096b7fa1063ca2fba19d98be6e6cdf" } +version = "0.3.0" +source = { git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.3.0#b309e2f01cda75dc96a6dc1a2355a7b3b64b5e16" } dependencies = [ { name = "absl-py" }, { name = "torch", marker = "sys_platform == 'never'" }, @@ -2279,7 +2279,7 @@ requires-dist = [ { name = "causal-conv1d", marker = "extra == 'ssm'", specifier = "~=1.5" }, { name = "datasets", marker = "extra == 'dev'" }, { name = "einops", marker = "extra == 'dev'", specifier = "~=0.8" }, - { name = "emerging-optimizers", marker = "extra == 'dev'", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.2.0" }, + { name = "emerging-optimizers", marker = "extra == 'dev'", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.3.0" }, { name = "fast-hadamard-transform", marker = "extra == 'dev'", git = "https://github.com/Dao-AILab/fast-hadamard-transform.git?rev=f134af63deb2df17e1171a9ec1ea4a7d8604d5ca" }, { name = "fastapi", marker = "extra == 'dev'", specifier = "~=0.50" }, { name = "flash-linear-attention", marker = "extra == 'dev'", specifier = "~=0.4.0" }, @@ -2351,7 +2351,7 @@ linting = [ { name = "ruff", specifier = "~=0.9.0" }, ] no-pypi-wheels = [ - { name = "emerging-optimizers", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.2.0" }, + { name = "emerging-optimizers", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.3.0" }, { name = "flash-mla", git = "https://github.com/deepseek-ai/FlashMLA?rev=nv_dev" }, ] test = [