[dev] fix no_shard training convergency and add unittest for no_shard

megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py

@@ -1242,6 +1242,9 @@ def start_param_sync(self, *unused, force_sync: bool = False, force_dispatch: bo
         """
         self._replace_param_with_raw_if_needed()
 
+        if self.data_parallel_sharding_strategy == "no_shard":
+            return
+
         if not force_sync and self.ddp_config.overlap_param_gather:
             # All-gather the first bucket before the forward pass.
             if self.ddp_config.fsdp_all_gather_in_start_param_sync:

megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py

@@ -3252,7 +3252,7 @@ def all_reduce_gradients(self, async_op: bool = False):
         all_reduce_ops = []
         for g in self.parameter_groups:
             gbuf = g.main_grad_buffer
-            if gbuf is not None:
+            if gbuf is None:
                 continue
             scaling_factor = gbuf.gradient_scaling_factor
             if self.ddp_config.check_for_nan_in_grad:

megatron/core/optimizer/__init__.py

@@ -939,6 +939,13 @@ def get_megatron_optimizer(
     model_chunk_offset = 0
     ddp_config = model_chunks[0].ddp_config  # Use the first model chunk's DDP config
     if ddp_config.use_megatron_fsdp:
+        # For no_shard, gradients are replicated across DP ranks after all-reduce, so grad stats
+        # should only be reduced over TP/PP (model_parallel_group) to avoid inflating the norm.
+        effective_intra_dist_opt_group = (
+            mp_group
+            if ddp_config.data_parallel_sharding_strategy == 'no_shard'
+            else intra_dist_opt_group
+        )
         for model_chunk, overlap_param_gather_with_optimizer_step in zip(
             all_dense_model_chunks, overlap_param_gather_with_optimizer_step_flags
         ):
@@ -960,7 +967,7 @@ def get_megatron_optimizer(
                 data_parallel_group=dp_cp_group,
                 data_parallel_group_gloo=intra_dp_cp_group_gloo,
                 data_parallel_group_idx=model_parallel_rank,
-                intra_dist_opt_group=intra_dist_opt_group,
+                intra_dist_opt_group=effective_intra_dist_opt_group,
                 distributed_optimizer_instance_id=distributed_optimizer_instance_id,
                 pg_collection=pg_collection,
             )

megatron/training/arguments.py

@@ -1274,6 +1274,12 @@ def validate_args(args, defaults={}):
         args.fsdp_manual_registration = True
         warn_rank_0('FSDP manual registration is enabled by default when nccl-ub is enabled')
 
+        if args.init_model_with_meta_device and args.data_parallel_sharding_strategy == "no_shard":
+            raise ValueError(
+                "Meta device initialization (init_model_with_meta_device=True) is not "
+                "supported or necessary for the 'no_shard' / 0 sharding strategy."
+            )
+
     if args.fsdp_manual_registration:
         assert (
             args.use_megatron_fsdp

tests/unit_tests/distributed/megatron_fsdp/test_mcore_fully_sharded_data_parallel.py

@@ -1,5 +1,6 @@
 # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
 import copy
+import gc
 import random
 
 import numpy as np
@@ -323,6 +324,10 @@ def train_step(model, optimizer, inputs):
                 msg=f"Parameters for {name1} don't match",
             )
 
+        gc.collect()
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+
     def test_fsdp_expt_device_mesh(self):
         """Test that expt_device_mesh is None for dense models and not None for MoE models."""
         if not is_torch_min_version("2.4.0"):
@@ -534,6 +539,10 @@ def train_step(model, optimizer, inputs):
                 msg=f"Parameters for {name1} don't match",
             )
 
+        gc.collect()
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+
     @classmethod
     def hsdp_one_step_test(cls, num_fsdp_group):
         if not is_torch_min_version("2.4.0"):
@@ -856,6 +865,10 @@ def _training_loop(seed=42, **kwargs):
                 dict(data_parallel_sharding_strategy="optim", fsdp_double_buffer=False),
                 id="optim_double_buffer",
             ),
+            pytest.param(
+                dict(data_parallel_sharding_strategy="no_shard", fsdp_double_buffer=False),
+                id="no_shard",
+            ),
         ],
     )
     def test_compatible_with_nd_parallel(self, ref_cache, nd_topology, spec_configs):
@@ -872,9 +885,13 @@ def test_compatible_with_nd_parallel(self, ref_cache, nd_topology, spec_configs)
                 use_distributed_optimizer=True, **distopt_spec_configs
             )
 
+        fsdp_sharding_strategy = spec_configs["data_parallel_sharding_strategy"]
+        # no_shard is incompatible with meta device initialization. See fully_shard.py:326.
+        init_model_with_meta_device = fsdp_sharding_strategy != "no_shard"
+
         outputs = TestMegatronFSDPE2E._training_loop(
             use_megatron_fsdp=True,
-            init_model_with_meta_device=True,
+            init_model_with_meta_device=init_model_with_meta_device,
             ckpt_format="fsdp_dtensor",
             gradient_accumulation_fusion=False,
             **spec_configs,
@@ -897,6 +914,10 @@ def test_compatible_with_nd_parallel(self, ref_cache, nd_topology, spec_configs)
                     ),
                 )
 
+        gc.collect()
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+
 
 def compare_losses(loss_a: float, loss_b: float, reference: str = "b"):
     """