diff --git a/megatron/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py
index f18d09f84e7..9ae23bb4b7f 100644
--- a/megatron/core/optimizer/optimizer.py
+++ b/megatron/core/optimizer/optimizer.py
@@ -1377,7 +1377,7 @@ def step(self):
                         or (
                             # Megatron-FSDP always uses decoupled_grad with FusedAdam.
                             self.config.use_precision_aware_optimizer
-                            and getattr(params[0], "__fsdp_param__", False)
+                            and getattr(parameters[0], "__fsdp_param__", False)
                         )
                     ),
                 )
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 62ef035dc7d..04f7c5c6482 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
@@ -773,7 +773,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"
     )
@@ -816,6 +815,7 @@ def _training_loop(seed=42, **kwargs):
                     data_parallel_sharding_strategy="optim_grads_params",
                     megatron_fsdp_main_params_dtype=torch.float32,
                     use_precision_aware_optimizer=True,
+                    fp8="hybrid",
                     fp8_recipe="delayed",
                     fp8_param_gather=True,
                     bf16=True,