[MXFP8/FP4-param-gather] Post processing after forced param AG in eval

megatron/core/distributed/distributed_data_parallel.py

@@ -7,7 +7,6 @@
 import torch
 
 from ..config_logger import has_config_logger_enabled, log_config_to_disk
-from ..fp8_utils import is_float8tensor, post_all_gather_processing
 from ..optimizer.param_layout import FullParamLayout
 from ..process_groups_config import ProcessGroupCollection
 from ..transformer.cuda_graphs import is_graph_capturing
@@ -476,53 +475,20 @@ def no_sync(self):
     def _start_bucket_group_param_sync(
         self, bucket_group: '_ParamAndGradBucketGroup', force_sync: bool
     ) -> None:
-        """Dispatch one bucket group's param all-gather + run the FP8 / MXFP8
+        """Dispatch one bucket group's param all-gather + run the FP8 / MXFP8 / FP4
         post-all-gather work the synchronous path needs.
 
         Factored out of :meth:`start_param_sync` so callers that own a subset
         of bucket groups (e.g. a chained ``LayerWiseDistributedOptimizer`` +
         ``DistributedOptimizer`` pair) can sync only their own buckets without
-        losing the FP8 post-processing that follows the collective.
+        losing the post-processing that follows the collective.
         """
         bucket_group.start_param_sync(force_sync=force_sync)
 
         if self.ddp_config.overlap_param_gather:
             return
 
-        # For MXFP8 params, we need to copy the all-gathered param data from the buffer to
-        # the param.data, since param buffer is not mapped to model params for MXFP8 case.
-        # The paramaters are cast from bf16 to MXFP8 during copy.
-        # In the case of "overlap_param_gather=True", the param copy is done
-        # in "finish_param_sync" stage after zeroing the shared gardient buffers.
-        if self.ddp_config.reuse_grad_buf_for_mxfp8_param_ag:
-            for bucket in bucket_group.buckets:
-                is_bf16_weight_bucket = False
-                for param in bucket.params:
-                    # Skip copying since bf16 weights in the mxfp8 model
-                    # are already mapped to param.data.
-                    if not is_float8tensor(param):
-                        is_bf16_weight_bucket = True
-                        break
-                    param_start, param_end = bucket.param_to_index[param]
-                    param_slice = bucket.param_data.view(-1)[param_start:param_end]
-                    param.data.copy_(param_slice.view(param.data.shape))
-                if is_bf16_weight_bucket:
-                    continue
-                # All-gathered params are not needed after being copied to param.data.
-                # Zero out the param buffer (shared with grad buffer) for gradient
-                # accumulation. We cannot zero out the entire grad buffer because one grad
-                # buffer may correspond to multiple param buffers. If we zero out the entire
-                # grad buffer, it would clear the data of those param buffers that have not
-                # yet completed AG.
-                bucket.param_data.zero_()
-        else:
-            fp8_params = []
-            for bucket in bucket_group.buckets:
-                for param in bucket.params:
-                    if is_float8tensor(param):
-                        fp8_params.append(param)
-            if len(fp8_params) > 0:
-                post_all_gather_processing(fp8_params)
+        bucket_group._post_param_sync()
 
     def start_param_sync(self, *unused, force_sync: bool = False, force_dispatch: bool = False):
         """

megatron/core/distributed/param_and_grad_buffer.py

@@ -267,6 +267,38 @@ def reset(self):
         self.per_param_grad_ready_counts = {}
         self.is_last_microbatch = True
 
+    def _post_param_sync(self):
+        """Run post-processing after param all-gather completes."""
+        if self.ddp_config.reuse_grad_buf_for_mxfp8_param_ag:
+            for bucket in self.buckets:
+                is_bf16_weight_bucket = False
+                for param in bucket.params:
+                    # Skip copying since bf16 weights in the mxfp8 model
+                    # are already mapped to param.data.
+                    if not is_float8tensor(param):
+                        is_bf16_weight_bucket = True
+                        break
+                    param_start, param_end = bucket.param_to_index[param]
+                    param_slice = bucket.param_data.view(-1)[param_start:param_end]
+                    param.data.copy_(param_slice.view(param.data.shape))
+                if is_bf16_weight_bucket:
+                    continue
+                # All-gathered params are not needed after being copied to param.data.
+                # Zero out the param buffer (shared with grad buffer) for gradient accumulation.
+                # We cannot zero out the entire grad buffer because one grad buffer may
+                # correspond to multiple param buffers. If we zero out the entire grad buffer,
+                # it would clear the data of those param buffers that have not yet completed AG.
+                bucket.param_data.zero_()
+            return
+
+        quantized_params = []
+        for bucket in self.buckets:
+            for param in bucket.params:
+                if is_float8tensor(param) or is_nvfp4tensor(param):
+                    quantized_params.append(param)
+        if len(quantized_params) > 0:
+            post_all_gather_processing(quantized_params)
+
     def check_grads(self, check_for_nan_or_inf, check_for_large):
         """
         Make sure norm of grads in bucket are not NaN prior to data-parallel
@@ -325,6 +357,7 @@ def start_param_sync(self, force_sync: bool = False):
             if self.param_gather_handle is not None:
                 self.param_gather_handle.wait()
                 self.param_gather_handle = None
+                self._post_param_sync()
                 return
         else:
             assert self.param_gather_handle is None
@@ -344,6 +377,8 @@ def start_param_sync(self, force_sync: bool = False):
             dp_size = self.intra_distributed_optimizer_instance_size
             if dp_size == 1:
                 # Single-rank group (e.g., expt_dp_size == 1): no all-gather needed.
+                if force_sync and self.ddp_config.overlap_param_gather:
+                    self._post_param_sync()
                 self.param_gather_dispatched = True
                 return
             local_rank = self.intra_distributed_optimizer_instance_rank
@@ -441,6 +476,8 @@ def start_param_sync(self, force_sync: bool = False):
                 # (async_op=False) is used, `cm` is not None. Manually set to None for
                 # consistency with prior code.
                 self.param_gather_handle = None
+        if force_sync and self.ddp_config.overlap_param_gather:
+            self._post_param_sync()
         self.param_gather_dispatched = True
 
     def finish_param_sync(self, skip_next_bucket_dispatch: bool = False):
@@ -480,30 +517,7 @@ def finish_param_sync(self, skip_next_bucket_dispatch: bool = False):
                 else:
                     self.next_param_gather_bucket_group.start_param_sync()
 
-            # For the mxfp8_param with "reuse_grad_buf_for_mxfp8_param_ag=True",
-            # we need to copy the param_data from the shared_param/grad_buffer to param.data
-            # after the param all-gather.
-            if self.ddp_config.reuse_grad_buf_for_mxfp8_param_ag:
-                for bucket in self.buckets:
-                    is_bf16_weight_bucket = False
-                    for param in bucket.params:
-                        # Skip copying since bf16 weights in the mxfp8 model
-                        # are already mapped to param.data.
-                        if not is_float8tensor(param):
-                            is_bf16_weight_bucket = True
-                            break
-                        param_start, param_end = bucket.param_to_index[param]
-                        param_slice = bucket.param_data.view(-1)[param_start:param_end]
-                        param.data.copy_(param_slice.view(param.data.shape))
-                    if is_bf16_weight_bucket:
-                        continue
-                    # All-gathered params are not needed after being copied to param.data.
-                    # Zero out the param buffer (shared with grad buffer) for gradient accumulation.
-                    # We cannot zero out the entire grad buffer because one grad buffer may
-                    # correspond to multiple param buffers. If we zero out the entire grad buffer,
-                    # it would clear the data of those param buffers that have not yet completed AG.
-                    bucket.param_data.zero_()
-            elif not self.ddp_config.use_distributed_optimizer:
+            if not self.ddp_config.use_distributed_optimizer:
                 for bucket in self.buckets:
                     if bucket.layerwise_gather_list is None:
                         continue
@@ -526,14 +540,7 @@ def finish_param_sync(self, skip_next_bucket_dispatch: bool = False):
                     # (a view into grad_data) would start from the result of the
                     # latest parameter all-gather instead of zero.
                     bucket.grad_data.zero_()
-            else:
-                fp8_params = []
-                for bucket in self.buckets:
-                    for param in bucket.params:
-                        if is_float8tensor(param):
-                            fp8_params.append(param)
-                if len(fp8_params) > 0:
-                    post_all_gather_processing(fp8_params)
+            self._post_param_sync()
 
     def start_grad_sync(self, force_all_reduce: Optional[bool] = False):
         """

megatron/training/training.py

@@ -2187,10 +2187,16 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch
         # 1. The first iteration's params are already in param.data (from init or checkpoint).
         # 2. Without forward_pre_hook, finish_param_sync() won't be called to zero the grad buffer,
         #    so the main grads will be polluted by the main params.
+        #
+        # Exception: when a full-iteration CUDA graph has been captured, the all-gather
+        # and subsequent param_data zero are baked into the graph and replay
+        # unconditionally. We must populate param_data so the replayed AG gathers
+        # correct weights, even when forward pre-hooks are disabled.
         if args.reuse_grad_buf_for_mxfp8_param_ag and args.overlap_param_gather:
             # Check if forward_pre_hook is enabled by checking if hooks are registered.
             forward_pre_hook_enabled = len(model[0].remove_forward_pre_hook_handles) > 0
-            if forward_pre_hook_enabled:
+            full_cg_captured = FullCudaGraphWrapper.cuda_graph.get("training") is not None
+            if forward_pre_hook_enabled or full_cg_captured:
                 for optim_instance in optimizer.chained_optimizers:
                     if isinstance(optim_instance, DistributedOptimizer):
                         optim_instance._copy_main_params_to_param_buffer()
@@ -3646,6 +3652,14 @@ def trace_handler(p):
             if args.log_energy:
                 energy_monitor.pause()
             timers('interval-time').stop()
+            if args.reuse_grad_buf_for_mxfp8_param_ag and args.overlap_param_gather:
+                # disable_forward_pre_hook(param_sync=True) below force-syncs params for eval.
+                # Copy the main params to param buffer before the forced AllGather.
+                for model_chunk in model:
+                    model_chunk.zero_grad_buffer()
+                for optim_instance in optimizer.chained_optimizers:
+                    if isinstance(optim_instance, DistributedOptimizer):
+                        optim_instance._copy_main_params_to_param_buffer()
             if should_disable_forward_pre_hook(args):
                 disable_forward_pre_hook(model)
                 pre_hook_enabled = False

tests/unit_tests/test_fp4_param.py

@@ -162,8 +162,37 @@ def get_batch(self, seq_length, micro_batch_size):
         loss_mask = torch.ones(seq_length).repeat((micro_batch_size, 1)).cuda()
         return input_ids, labels, position_ids, attention_mask, loss_mask
 
+    def run_eval_transition(self, args, model_chunks, batch):
+        input_ids, labels, position_ids, attention_mask, loss_mask = batch
+
+        if should_disable_forward_pre_hook(args):
+            disable_forward_pre_hook(model_chunks, param_sync=True)
+
+        model_chunks[0].eval()
+        model_chunks[0].set_is_first_microbatch()
+        with torch.no_grad():
+            eval_output = model_chunks[0].forward(
+                input_ids=input_ids,
+                position_ids=position_ids,
+                attention_mask=attention_mask,
+                labels=labels,
+                loss_mask=loss_mask,
+            )
+        eval_loss = eval_output.mean()
+        model_chunks[0].train()
+
+        if should_disable_forward_pre_hook(args):
+            enable_forward_pre_hook(model_chunks)
+
+        return eval_loss.item()
+
     def _run_test_helper(
-        self, tp_size, inference: bool = False, fp4_param_gather: bool = True, **kwargs
+        self,
+        tp_size,
+        inference: bool = False,
+        fp4_param_gather: bool = True,
+        eval_transition: bool = False,
+        **kwargs,
     ):
         """Test fp4_param with gpt_model."""
         args = self.create_test_args(
@@ -206,6 +235,7 @@ def _run_test_helper(
             assert num_fp4_params == 4 * fp4_layers
 
         loss_list = []
+        eval_loss_list = []
 
         # CUDA graph setup (transformer_engine implementation)
         cuda_graph_helper = None
@@ -267,6 +297,17 @@ def _run_test_helper(
 
             loss_list.append(loss.item())
 
+            if eval_transition:
+                eval_loss_list.append(
+                    self.run_eval_transition(
+                        args,
+                        gpt_model,
+                        (input_ids, labels, position_ids, attention_mask, loss_mask),
+                    )
+                )
+
+        if eval_transition:
+            return torch.tensor(loss_list), torch.tensor(eval_loss_list)
         return torch.tensor(loss_list)
 
     def run_test(self, tp_size, inference: bool = False, **kwargs):
@@ -282,6 +323,18 @@ def run_test(self, tp_size, inference: bool = False, **kwargs):
 
             torch.testing.assert_close(loss_list, loss_list_ref, atol=1e-2, rtol=1e-2)
 
+    def run_test_with_eval_transition(self, tp_size, **kwargs):
+        """Test fp4_param eval transition with gpt_model."""
+        loss_list, eval_loss_list = self._run_test_helper(
+            tp_size, fp4_param_gather=True, eval_transition=True, **kwargs
+        )
+        loss_list_ref, eval_loss_list_ref = self._run_test_helper(
+            tp_size, fp4_param_gather=False, eval_transition=True, **kwargs
+        )
+
+        torch.testing.assert_close(loss_list, loss_list_ref, atol=1e-2, rtol=1e-2)
+        torch.testing.assert_close(eval_loss_list, eval_loss_list_ref, atol=1e-2, rtol=1e-2)
+
     @pytest.mark.skipif(not is_nvfp4_available, reason=reason_for_no_nvfp4)
     @pytest.mark.skipif(not is_te_min_version("2.7.0.dev0"), reason="TE 2.7.0.dev0 is required")
     @pytest.mark.parametrize("tp_size", [2])
@@ -294,6 +347,13 @@ def test_nvfp4(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, inference=False, **kwargs)
 
+    @pytest.mark.skipif(not is_nvfp4_available, reason=reason_for_no_nvfp4)
+    @pytest.mark.skipif(not is_te_min_version("2.7.0.dev0"), reason="TE 2.7.0.dev0 is required")
+    @pytest.mark.parametrize("tp_size", [2])
+    def test_nvfp4_eval_transition(self, tp_size):
+        kwargs = {"overlap_param_gather": True, "overlap_grad_reduce": True}
+        self.run_test_with_eval_transition(tp_size=tp_size, **kwargs)
+
     @pytest.mark.skipif(not is_nvfp4_available, reason=reason_for_no_nvfp4)
     @pytest.mark.skipif(not is_te_min_version("2.7.0.dev0"), reason="TE 2.7.0.dev0 is required")
     @pytest.mark.parametrize("tp_size", [2])