[sync] DeepSeek V3 FSDP Fix for Precision-Aware Optimizer

src/megatron/bridge/training/config.py

@@ -1427,9 +1427,6 @@ def validate(self) -> None:
                 print_rank_0("average_in_collective is not supported with Megatron FSDP, setting to True")
                 self.ddp.average_in_collective = False
 
-            if self.optimizer.use_precision_aware_optimizer:
-                self.ddp.preserve_fp32_weights = False
-
             # TODO: This can be removed once NVIDIA/TransformerEngine#2371 is available to use
             if self.model.gradient_accumulation_fusion:
                 print_rank_0("Gradient accumulation fusion is not supported with Megatron FSDP, setting to False")

tests/unit_tests/training/test_config.py

@@ -1564,70 +1564,6 @@ def test_fsdp_dtensor_format_validation_without_megatron_fsdp_fails(self, monkey
         finally:
             restore_get_world_size_safe(og_ws, cfg_mod)
 
-    def test_megatron_fsdp_with_precision_aware_optimizer(self, monkeypatch):
-        """Test that Megatron FSDP with precision aware optimizer sets preserve_fp32_weights=False."""
-        gpt_model_cfg = create_test_gpt_config()
-        train_cfg = create_test_training_config(train_iters=500, global_batch_size=16)
-        sched_cfg = create_test_scheduler_config()
-
-        # Create optimizer config with precision aware optimizer enabled
-        optim_cfg = create_test_optimizer_config()
-        optim_cfg.use_precision_aware_optimizer = True
-        optim_cfg.use_distributed_optimizer = True  # Required for precision aware optimizer
-
-        dist_cfg = create_test_distributed_init_config(use_megatron_fsdp=True)
-        ddp_cfg = create_test_ddp_config(use_distributed_optimizer=True)
-
-        container, og_ws, cfg_mod = create_test_config_container(
-            world_size_override=1,
-            model_config=gpt_model_cfg,
-            train_config=train_cfg,
-            scheduler_config=sched_cfg,
-            optimizer_config=optim_cfg,
-            dist_config=dist_cfg,
-            ddp_config=ddp_cfg,
-        )
-        try:
-            container.validate()
-            # Should automatically set preserve_fp32_weights=False when using precision aware optimizer with FSDP
-            assert container.ddp.preserve_fp32_weights is False
-            assert container.optimizer.use_precision_aware_optimizer is True
-            assert container.dist.use_megatron_fsdp is True
-        finally:
-            restore_get_world_size_safe(og_ws, cfg_mod)
-
-    def test_megatron_fsdp_without_precision_aware_optimizer(self, monkeypatch):
-        """Test that Megatron FSDP without precision aware optimizer doesn't modify preserve_fp32_weights."""
-        gpt_model_cfg = create_test_gpt_config()
-        train_cfg = create_test_training_config(train_iters=500, global_batch_size=16)
-        sched_cfg = create_test_scheduler_config()
-
-        # Create optimizer config with precision aware optimizer disabled
-        optim_cfg = create_test_optimizer_config()
-        optim_cfg.use_precision_aware_optimizer = False
-        optim_cfg.use_distributed_optimizer = True  # Enable distributed optimizer for consistency
-
-        dist_cfg = create_test_distributed_init_config(use_megatron_fsdp=True)
-        ddp_cfg = create_test_ddp_config(use_distributed_optimizer=True)
-
-        container, og_ws, cfg_mod = create_test_config_container(
-            world_size_override=1,
-            model_config=gpt_model_cfg,
-            train_config=train_cfg,
-            scheduler_config=sched_cfg,
-            optimizer_config=optim_cfg,
-            dist_config=dist_cfg,
-            ddp_config=ddp_cfg,
-        )
-        try:
-            container.validate()
-            # preserve_fp32_weights should keep its default value when precision aware optimizer is disabled
-            assert container.optimizer.use_precision_aware_optimizer is False
-            assert container.dist.use_megatron_fsdp is True
-            # preserve_fp32_weights should remain at its default
-        finally:
-            restore_get_world_size_safe(og_ws, cfg_mod)
-
 
 class TestMixedPrecisionConsistencyValidation:
     """Tests for _validate_mixed_precision_consistency function.