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 04f7c5c6482..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,12 +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( @@ -764,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