[Main] Fix FusedAdam.use_decoupled_grad mis-set for Megatron-FSDP.

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

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