Fix mxfp8 param gather numerical issue when DP overlap is off

megatron/core/optimizer/distrib_optimizer.py

@@ -3017,10 +3017,12 @@ def step_with_ready_grads(self) -> bool:
         """
         update_successful = super().step_with_ready_grads()
 
+        should_sync_params = not self.ddp_config.overlap_param_gather and not getattr(
+            self, '_defer_param_sync', False
+        )
         timers = self.config.timers
-        if timers is not None:
+        if timers is not None and (self.ddp_config.use_megatron_fsdp or should_sync_params):
             timers('params-all-gather', log_level=1).start(barrier=self.config.barrier_with_L1_time)
-
         if self.ddp_config.use_megatron_fsdp:
             # Optionally all-gather Megatron-FSDP sharded main weights
             # early in preparation for the subsequent forward pass.
@@ -3031,12 +3033,12 @@ def step_with_ready_grads(self) -> bool:
             # communication calls here. If overlapping all-gather for parameters, the following
             # the first all-gather is launched asynchronously in the next optimizer.zero_grad()
             # call and subsequent all-gathers are launched in the forward pre-hook.
-            if not self.ddp_config.overlap_param_gather:
+            if should_sync_params:
                 # Only sync DistOpt-managed bucket groups so a sibling
                 # LayerWiseDistributedOptimizer's own ``start_param_sync`` call
                 # is not duplicated for the same buckets.
                 self.start_param_sync_for_bucket_group_subset()
-        if timers is not None:
+        if timers is not None and (self.ddp_config.use_megatron_fsdp or should_sync_params):
             timers('params-all-gather').stop()
 
         return update_successful

megatron/core/optimizer/optimizer.py

@@ -1280,19 +1280,100 @@ def prepare_grads(self) -> bool:
 
         return found_inf_flag
 
-    @torch.no_grad()
-    def step_with_ready_grads(self) -> bool:
-        """Step the optimizer with ready gradients, return successful."""
+    def _step(self) -> bool:
+        """Step all optimizers in this chain."""
         success = True
         for optimizer_idx, optimizer in enumerate(self.chained_optimizers):
             success &= optimizer.step_with_ready_grads()
             if self.config.overlap_param_gather_with_optimizer_step and optimizer_idx == 0:
                 assert success
                 assert len(optimizer.model_chunks) == 1
                 optimizer.model_chunks[0].start_param_sync(force_dispatch=True)
+        return success
+
+    def _should_defer_mxfp8_param_sync(self) -> bool:
+        """Return whether MXFP8 param sync should be deferred until chained steps finish."""
+        return (
+            self.config.reuse_grad_buf_for_mxfp8_param_ag and not self.config.overlap_param_gather
+        )
+
+    def _enable_deferred_mxfp8_param_sync(self) -> List[Tuple[Any, Any]]:
+        """Enable deferred DistOpt param sync and collect bucket groups to sync later."""
+        from .distrib_optimizer import DistributedOptimizer
+        from .layer_wise_optimizer import _bucket_is_managed_by_layer_wise_optimizer
+
+        # With MXFP8 grad-buffer reuse and non-overlap param gather, each DistOpt stages
+        # its own updated main-param shards into its param buffers during step. However,
+        # param sync is a DDP bucket-group operation that copies gathered values into model
+        # weights and zeros the shared MXFP8 param/grad buffers. For MoE, dense and expert
+        # DistOpts may share the same model chunk, so defer param sync until all chained
+        # optimizers have staged their params, then sync each DistOpt-managed bucket group once.
+        deferred_bucket_groups = []
+        deferred_bucket_group_ids = set()
+
+        for optimizer in self.chained_optimizers:
+            if not isinstance(optimizer, DistributedOptimizer):
+                continue
+
+            optimizer._defer_param_sync = True
+            for model_chunk in optimizer.model_chunks:
+                for bucket_group in (
+                    model_chunk.bucket_groups + model_chunk.expert_parallel_bucket_groups
+                ):
+                    if not bucket_group.buckets:
+                        continue
+                    if _bucket_is_managed_by_layer_wise_optimizer(
+                        bucket_group.buckets[0], default_for_untagged=False
+                    ):
+                        continue
+
+                    bucket_group_id = id(bucket_group)
+                    if bucket_group_id in deferred_bucket_group_ids:
+                        continue
+
+                    deferred_bucket_group_ids.add(bucket_group_id)
+                    deferred_bucket_groups.append((model_chunk, bucket_group))
+
+        return deferred_bucket_groups
+
+    def _disable_deferred_mxfp8_param_sync(self) -> None:
+        """Disable deferred DistOpt param sync."""
+        for optimizer in self.chained_optimizers:
+            if hasattr(optimizer, '_defer_param_sync'):
+                optimizer._defer_param_sync = False
+
+    def _start_deferred_mxfp8_param_sync(
+        self, deferred_bucket_groups: List[Tuple[Any, Any]]
+    ) -> None:
+        """Start param sync for deferred bucket groups."""
+        timers = self.config.timers
+        if timers is not None:
+            timers('params-all-gather', log_level=1).start(barrier=self.config.barrier_with_L1_time)
+        for model_chunk, bucket_group in deferred_bucket_groups:
+            model_chunk._start_bucket_group_param_sync(bucket_group, force_sync=False)
+        if timers is not None:
+            timers('params-all-gather').stop()
+
+    def _step_with_deferred_mxfp8_param_sync(self) -> bool:
+        """Step optimizers with MXFP8 param sync deferred until all steps finish."""
+        deferred_bucket_groups = self._enable_deferred_mxfp8_param_sync()
+        try:
+            success = self._step()
+        finally:
+            self._disable_deferred_mxfp8_param_sync()
+
+        if success and deferred_bucket_groups:
+            self._start_deferred_mxfp8_param_sync(deferred_bucket_groups)
 
         return success
 
+    @torch.no_grad()
+    def step_with_ready_grads(self) -> bool:
+        """Step the optimizer with ready gradients, return successful."""
+        if self._should_defer_mxfp8_param_sync():
+            return self._step_with_deferred_mxfp8_param_sync()
+        return self._step()
+
     def grads_states_parallel_group_is_shared(self):
         """Check if all optimizers share the same gradient statistics parallel group."""
         reference_group = self.chained_optimizers[0].get_grad_stats_parallel_group()

tests/unit_tests/test_fp8_param.py

@@ -90,7 +90,9 @@ def model_provider(
         model_parallel_cuda_manual_seed(_SEED)
         args = get_args()
         config = core_transformer_config_from_args(args)
-        transformer_layer_spec = layer_spec_fn()
+        transformer_layer_spec = layer_spec_fn(
+            num_experts=args.num_experts, moe_grouped_gemm=args.moe_grouped_gemm
+        )
         return GPTModel(
             config=config,
             transformer_layer_spec=transformer_layer_spec,
@@ -180,7 +182,7 @@ def _run_test_helper(
         use_cuda_graph: bool = False,
         **kwargs,
     ):
-        """Test fp8_param with gpt_model."""
+        """Test fp8_param with a small GPT model."""
         args = self.create_test_args(
             tp_size,
             recipe,
@@ -199,7 +201,10 @@ def _run_test_helper(
 
         set_args(args)
         torch.manual_seed(_SEED)
-        Utils.initialize_model_parallel(tensor_model_parallel_size=tp_size)
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=tp_size,
+            expert_model_parallel_size=args.expert_model_parallel_size,
+        )
 
         input_ids, labels, position_ids, attention_mask, loss_mask = self.get_batch(
             self.seq_length, self.micro_batch_size
@@ -237,14 +242,21 @@ def _run_test_helper(
             if is_float8tensor(param):
                 num_fp8_params += 1
 
-        # Verify the number of fp8 params.
         fp8_layers = args.num_layers
         if kwargs.get("first_last_layers_bf16", False):
             fp8_layers -= kwargs["num_layers_at_start_in_bf16"]
             fp8_layers -= kwargs["num_layers_at_end_in_bf16"]
-        # Each layer has 4 GEMM weights: qkv, proj, fc1, fc2.
-        if fp8_param_gather:
-            assert num_fp8_params == 4 * fp8_layers
+        if fp8_param_gather and fp8_layers > 0:
+            if args.num_experts is None:
+                # Each dense layer has 4 GEMM weights: qkv, proj, fc1, fc2.
+                assert num_fp8_params == 4 * fp8_layers
+            else:
+                assert num_fp8_params > 0
+                assert any(
+                    not getattr(param, 'allreduce', True) for param in gpt_model[0].parameters()
+                )
+                if not inference:
+                    assert len(optimizer.chained_optimizers) >= 2
 
         # Verify that bf16 params (embedding, LN, etc.) in the MXFP8 model are mapped
         # to the param buffer (shared with grad buffer) rather than allocated separately.
@@ -332,7 +344,7 @@ def _run_test_helper(
         return torch.tensor(loss_list)
 
     def run_test(self, tp_size, recipe, inference: bool = False, **kwargs):
-        """Test fp8_param with gpt_model."""
+        """Test fp8_param with a small GPT model."""
         if inference:
             with torch.inference_mode():
                 self._run_test_helper(tp_size, recipe, inference=True, **kwargs)
@@ -442,6 +454,36 @@ def test_mxfp8(self, tp_size, dp_overlap):
         kwargs = {"overlap_param_gather": dp_overlap[0], "overlap_grad_reduce": dp_overlap[1]}
         self.run_test(tp_size=tp_size, recipe="mxfp8", **kwargs)
 
+    @pytest.mark.skipif(
+        get_device_arch_version() < 10, reason="MXFP8 is supported since Blackwell architecture"
+    )
+    @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+    @pytest.mark.skipif(not is_te_min_version("2.3.0.dev0"), reason="TE 2.3.0.dev0 is required")
+    @pytest.mark.parametrize("tp_size", [1])
+    @pytest.mark.parametrize("dp_overlap", [(False, False), (False, True), (True, True)])
+    def test_mxfp8_moe(self, tp_size, dp_overlap):
+        """
+        dp_overlap: (overlap_param_gather, overlap_grad_reduce)
+        """
+        kwargs = {
+            "overlap_param_gather": dp_overlap[0],
+            "overlap_grad_reduce": dp_overlap[1],
+            "num_layers": 4,
+            "vocal_size": 128800,
+            "hidden_size": 128,
+            "num_attention_heads": 8,
+            "expert_model_parallel_size": 2,
+            "num_experts": 2,
+            "moe_grouped_gemm": True,
+            "moe_token_dispatcher_type": "alltoall",
+            "moe_router_topk": 1,
+            "moe_router_pre_softmax": True,
+            "moe_router_load_balancing_type": "none",
+            "moe_aux_loss_coeff": 0.0,
+            "moe_ffn_hidden_size": 128,
+        }
+        self.run_test(tp_size=tp_size, recipe="mxfp8", **kwargs)
+
     @pytest.mark.skipif(
         get_device_arch_version() < 10, reason="MXFP8 is supported since Blackwell architecture"
     )