From 5b88b399d50ee567bccd92fb14aa7864a55454e6 Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Wed, 22 Apr 2026 08:20:54 -0700 Subject: [PATCH 1/2] Fix segfault caused by not using decoupled gradient for Megatron-FSDP. Signed-off-by: Cory Ye --- megatron/core/optimizer/__init__.py | 2 ++ .../megatron_fsdp/test_mcore_fully_sharded_data_parallel.py | 4 +++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py index 4f69a9efd55..c6d3e41aed5 100644 --- a/megatron/core/optimizer/__init__.py +++ b/megatron/core/optimizer/__init__.py @@ -975,6 +975,8 @@ def get_megatron_optimizer( # applied to the Megatron-FSDP main weight and extended to FusedAdam # main weights. Override this here. setattr(optimizer_part.optimizer, "master_weights", False) + # Megatron-FSDP always uses a decoupled gradient when using FusedAdam. + setattr(optimizer_part.optimizer, "use_decoupled_grad", True) optimizers.append(optimizer_part) model_chunk_offset += 1 diff --git a/tests/unit_tests/distributed/megatron_fsdp/test_mcore_fully_sharded_data_parallel.py b/tests/unit_tests/distributed/megatron_fsdp/test_mcore_fully_sharded_data_parallel.py index 802b3577d9d..a2b6a18680c 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/test_mcore_fully_sharded_data_parallel.py +++ b/tests/unit_tests/distributed/megatron_fsdp/test_mcore_fully_sharded_data_parallel.py @@ -742,6 +742,9 @@ def _training_loop(seed=42, **kwargs): assert ( not optim.optimizer.master_weights ), "Megatron-FSDP should not use FusedAdam master weights." + assert ( + optim.optimizer.use_decoupled_grad + ), "Megatron-FSDP should be using a decoupled gradient with FusedAdam." # Prepare data iterator data_iterator = make_gpt_mock_data_iterator( @@ -773,7 +776,6 @@ def _training_loop(seed=42, **kwargs): return outputs - @pytest.mark.flaky_in_dev @pytest.mark.skipif( not is_torch_min_version("2.4.0"), reason="Test needs to be updated for torch >= 2.4.0" ) From aba018ec9ee22917beca2c79e731fe0a1c4cb77f Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Thu, 23 Apr 2026 20:24:17 -0700 Subject: [PATCH 2/2] Add more tests catching the mis-placed gradient and FusedAdam(use_decoupled_grad=False). Signed-off-by: Cory Ye --- .../test_mcore_fully_sharded_data_parallel.py | 21 +++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/tests/unit_tests/distributed/megatron_fsdp/test_mcore_fully_sharded_data_parallel.py b/tests/unit_tests/distributed/megatron_fsdp/test_mcore_fully_sharded_data_parallel.py index a2b6a18680c..500045871e7 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/test_mcore_fully_sharded_data_parallel.py +++ b/tests/unit_tests/distributed/megatron_fsdp/test_mcore_fully_sharded_data_parallel.py @@ -736,15 +736,21 @@ def _training_loop(seed=42, **kwargs): train_iters=NUM_TRAINING_STEPS, **kwargs, ) - if kwargs.get("use_megatron_fsdp", False) and kwargs.get( + megatron_fsdp_te_fused_adam = kwargs.get("use_megatron_fsdp", False) and kwargs.get( "use_precision_aware_optimizer", False - ): + ) + if megatron_fsdp_te_fused_adam: assert ( not optim.optimizer.master_weights ), "Megatron-FSDP should not use FusedAdam master weights." assert ( optim.optimizer.use_decoupled_grad ), "Megatron-FSDP should be using a decoupled gradient with FusedAdam." + assert model_chunks[ + 0 + ].module.param_and_grad_buffer.use_decoupled_grad, ( + "Megatron-FSDP is installing gradients into param.decoupled_grad." + ) # Prepare data iterator data_iterator = make_gpt_mock_data_iterator( @@ -767,6 +773,17 @@ def _training_loop(seed=42, **kwargs): micro_batch_size=MICRO_BATCH_SIZE, num_micro_batches=GLOBAL_BATCH_SIZE // MICRO_BATCH_SIZE // DP_GROUP.size(), ) + # Check that at least one non-null / non-zero gradient + # exists when using Megatron-FSDP. + if kwargs.get("use_megatron_fsdp", False): + grad_attr = "decoupled_grad" if megatron_fsdp_te_fused_adam else "grad" + assert any( + [ + getattr(p, grad_attr, None) is not None + and getattr(p, grad_attr, None)._local_tensor.any() + for p in model_chunks[0].parameters() + ] + ), f"[Megatron-FSDP] Missing gradient in Parameter.{grad_attr}..." optim.step() # Collect loss