From 4b1a3cc4bbcfacd7ddfdf5e8581ff264bfb5f3f5 Mon Sep 17 00:00:00 2001
From: Deepak Narayanan <dnarayanan@nvidia.com>
Date: Sun, 4 Jan 2026 08:43:39 -0800
Subject: [PATCH 1/5] Add ability to save wgrads and dgrads in a compatible way
 that shares code

wgrad saving requires that the DP collective is an all-reduce instead of reduce-scatter to make it easier to pull wgrads from DP replica 0. Some other gotchas:
- First gradient reduction (when metadata is being collected about when to launch collectives) should also be an all-reduce if needed
- start_grad_sync call should be a no-op if in first batch and collective has already been dispatched
- Don't try calling reduce_scatter_with_fp32_accumulation's .wait() method if using all-reduce in a particular iteration

Signed-off-by: Deepak Narayanan <dnarayanan@nvidia.com>
---
 .../distributed/distributed_data_parallel.py  |   9 +-
 .../core/distributed/finalize_model_grads.py  |   3 +-
 .../core/distributed/param_and_grad_buffer.py |  22 ++--
 megatron/core/pipeline_parallel/schedules.py  |  14 ++-
 megatron/training/arguments.py                |   4 +
 megatron/training/checkpointing.py            |  34 ++++++
 megatron/training/dgrad_logging.py            | 114 ++++++++++++++++++
 megatron/training/training.py                 |  41 ++++++-
 8 files changed, 224 insertions(+), 17 deletions(-)
 create mode 100644 megatron/training/dgrad_logging.py

diff --git a/megatron/core/distributed/distributed_data_parallel.py b/megatron/core/distributed/distributed_data_parallel.py
index e831d7cf4ec..421b279b17d 100644
--- a/megatron/core/distributed/distributed_data_parallel.py
+++ b/megatron/core/distributed/distributed_data_parallel.py
@@ -93,6 +93,7 @@ def __init__(
         # disable_bucketing is True (e.g., we might not want to break up model parameters
         # into buckets for model chunks after the first in the interleaved schedule).
         self.bucket_size = self.ddp_config.bucket_size
+        self.force_all_reduce = False
         if isinstance(self.pp_group, list):
             pp_rank = self.pp_group[0].rank()
         else:
@@ -440,7 +441,9 @@ def hook(*unused):
                 param.grad = None
 
                 if self.ddp_config.overlap_grad_reduce:
-                    self.param_to_bucket_group[param].register_grad_ready(param)
+                    self.param_to_bucket_group[param].register_grad_ready(
+                        param, self.force_all_reduce
+                    )
 
         return hook
 
@@ -519,7 +522,7 @@ def start_grad_sync(self, *unused):
         for bucket_group in self.bucket_groups + self.expert_parallel_bucket_groups:
             bucket_group.start_grad_sync()
 
-    def finish_grad_sync(self):
+    def finish_grad_sync(self, force_all_reduce: Optional[bool] = False):
         """
         Finishes grad sync (all-reduce or reduce-scatter) communication operations
         for all model gradients.
@@ -529,7 +532,7 @@ def finish_grad_sync(self):
         communication ops.
         """
         for bucket_group in self.bucket_groups + self.expert_parallel_bucket_groups:
-            bucket_group.finish_grad_sync()
+            bucket_group.finish_grad_sync(force_all_reduce=force_all_reduce)
 
     def scale_gradients(self, scaling_factor: float):
         """Scale all gradients inside the buffers by `scaling_factor`."""
diff --git a/megatron/core/distributed/finalize_model_grads.py b/megatron/core/distributed/finalize_model_grads.py
index ddaeb7e8d84..a52592bb269 100644
--- a/megatron/core/distributed/finalize_model_grads.py
+++ b/megatron/core/distributed/finalize_model_grads.py
@@ -397,6 +397,7 @@ def finalize_model_grads(
     model: List[torch.nn.Module],
     num_tokens: Optional[torch.Tensor] = None,
     pg_collection: Optional[ProcessGroupCollection] = None,
+    force_all_reduce: Optional[bool] = False,
 ):
     """
     All-reduce all model grads across DP replicas, layernorm grads for sequence parallelism,
@@ -439,7 +440,7 @@ def finalize_model_grads(
     if config.timers is not None:
         config.timers('all-grads-sync', log_level=1).start(barrier=config.barrier_with_L1_time)
     for model_chunk in model:
-        model_chunk.finish_grad_sync()
+        model_chunk.finish_grad_sync(force_all_reduce=force_all_reduce)
     if config.timers is not None:
         config.timers('all-grads-sync').stop()
 
diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py
index 50cf3e0ea37..e7e32ddc081 100644
--- a/megatron/core/distributed/param_and_grad_buffer.py
+++ b/megatron/core/distributed/param_and_grad_buffer.py
@@ -355,7 +355,7 @@ def finish_param_sync(self, skip_next_bucket_dispatch: bool = False):
                 if len(fp8_params) > 0:
                     post_all_gather_processing(fp8_params)
 
-    def start_grad_sync(self):
+    def start_grad_sync(self, force_all_reduce: Optional[bool] = False):
         """
         Initiates grad sync (all-reduce or reduce-scatter) communication operations
         for all buckets in the bucket group.
@@ -426,7 +426,7 @@ def start_grad_sync(self):
         grad_reduce_handle = None
         with stream_context, _coalescing_manager(communication_group, async_ops=async_op) as cm:
             for idx, bucket in enumerate(self.buckets):
-                if self.ddp_config.use_distributed_optimizer:
+                if self.ddp_config.use_distributed_optimizer and not force_all_reduce:
                     if self.cached_grad_buffer_shard_list[idx] is None:
                         self.cached_grad_buffer_shard_list[idx] = shard_buffer(
                             bucket.grad_data, self.intra_distributed_optimizer_instance_size
@@ -442,6 +442,10 @@ def start_grad_sync(self):
                         async_op=async_op,
                     )
                 else:
+                    if torch.distributed.get_rank() == 0 and force_all_reduce:
+                        logger.info(
+                            f"Performing reduction using all_reduce because {force_all_reduce=}"
+                        )
                     torch.distributed.all_reduce(
                         bucket.grad_data, op=reduce_op, group=communication_group, async_op=async_op
                     )
@@ -476,7 +480,7 @@ def start_grad_sync(self):
                     )
 
         if async_op:
-            if self.ddp_config.reduce_scatter_with_fp32_accumulation:
+            if self.ddp_config.reduce_scatter_with_fp32_accumulation and not force_all_reduce:
                 assert (
                     len(self.buckets) == 1
                 ), "Only 1 bucket supported with reduce_scatter_with_fp32_accumulation=True"
@@ -494,7 +498,7 @@ def start_grad_sync(self):
             # None.
             self.grad_reduce_handle = None
 
-    def finish_grad_sync(self):
+    def finish_grad_sync(self, force_all_reduce: Optional[bool] = False):
         """
         Finishes grad sync (all-reduce or reduce-scatter) communication operations
         for all buckets in the bucket group.
@@ -506,13 +510,13 @@ def finish_grad_sync(self):
         self.param_gather_dispatched = False
         # If overlap_grad_reduce is False, start (and finish) synchronous communication call here.
         if not self.ddp_config.overlap_grad_reduce:
-            self.start_grad_sync()
+            self.start_grad_sync(force_all_reduce=force_all_reduce)
             return
         # If first batch, start asynchronous communication here. register_grad_ready() launches
         # asynchronous communication only once self.golden_per_param_grad_ready_counts is
         # populated at the end of this first batch.
         if self.is_first_batch:
-            self.start_grad_sync()
+            self.start_grad_sync(force_all_reduce=force_all_reduce)
         # When using multiple DistOpt instances, we don't need to sync here as we launch
         # communications on a separate communication stream.
         if self.ddp_config.num_distributed_optimizer_instances > 1:
@@ -526,7 +530,9 @@ def finish_grad_sync(self):
         self.grad_reduce_handle.wait()
         self.grad_reduce_handle = None
 
-    def register_grad_ready(self, param: torch.nn.Parameter):
+    def register_grad_ready(
+        self, param: torch.nn.Parameter, force_all_reduce: Optional[bool] = False
+    ):
         """
         Registers grads for the passed-in param to be "ready" for grad sync.
 
@@ -546,7 +552,7 @@ def register_grad_ready(self, param: torch.nn.Parameter):
             if not self.is_first_batch:
                 if self.per_param_grad_ready_counts == self.golden_per_param_grad_ready_counts:
                     assert len(self.per_param_grad_ready_counts) == len(self.params)
-                    self.start_grad_sync()
+                    self.start_grad_sync(force_all_reduce=force_all_reduce)
 
 
 class _ParamAndGradBuffer:
diff --git a/megatron/core/pipeline_parallel/schedules.py b/megatron/core/pipeline_parallel/schedules.py
index f1bfc340faf..07b8c4c54e9 100644
--- a/megatron/core/pipeline_parallel/schedules.py
+++ b/megatron/core/pipeline_parallel/schedules.py
@@ -111,9 +111,9 @@ def forward_step(data_iterator, model):
     decoder_seq_length (int, optional): The sequence length for the decoder in a dual-stack
         transformer. This is ignored for a single-stack transformer.
 
-    forward_only (optional, default = False): Perform only the forward step
+    forward_only (optional, default = False): Perform only the forward step.
 
-    collect_non_loss_data (optional, bool, default=False): TODO
+    collect_non_loss_data (optional, bool, default=False): TODO.
 
     first_val_step (bool, optional): Is the first step of the validation phase. Used by
         Transformer Engine modules to only update their fp8 weights only on the first validation
@@ -125,11 +125,15 @@ def forward_step(data_iterator, model):
         respective list of shapes. Thus it is not used in the other forward-backward functions
         which have different shape handling.
 
+    force_all_reduce (bool, optional): Run all-reduce for this iteration instead of maybe
+        reduce-scatter to ensure all ranks have fully reduced gradients.
+
     Args:
         pp_size (Optional[int]): Pipeline model parallel size to use.
         vp_size (Optional[int]): Virtual pipeline model parallel size to use.
             If both pp_size and vp_size are None, both values fall back to parallel_state.
             Otherwise, provided values are used as-is and None is treated as an explicit input.
+
     """
     if pp_size is None and vp_size is None:
         pp_size = parallel_state.get_pipeline_model_parallel_world_size()
@@ -524,6 +528,7 @@ def forward_backward_no_pipelining(
     adjust_tensor_shapes_fn: Optional[Callable] = None,  # unused
     p2p_communicator: Optional[P2PCommunicator] = None,  # unused
     pg_collection: Optional[ProcessGroupCollection] = None,
+    force_all_reduce: Optional[bool] = False,
 ):
     """Run forward and backward passes with no pipeline parallelism"""
 
@@ -674,6 +679,7 @@ def forward_backward_no_pipelining(
             [model],
             total_num_tokens if config.calculate_per_token_loss else None,
             pg_collection=pg_collection,
+            force_all_reduce=force_all_reduce,
         )
 
     if not forward_only and config.fine_grained_activation_offloading:
@@ -832,6 +838,7 @@ def forward_backward_pipelining_with_interleaving(
     adjust_tensor_shapes_fn: Optional[Callable] = None,  # unused
     p2p_communicator: Optional[P2PCommunicator] = None,
     pg_collection: Optional[ProcessGroupCollection] = None,
+    force_all_reduce: Optional[bool] = False,
 ):
     """Run interleaved 1F1B schedule (model split into model chunks), with
     communication between pipeline stages as needed.
@@ -1911,6 +1918,7 @@ def pp_post_backward(input_tensor_grad, vp_stage=None):
             model,
             total_num_tokens if config.calculate_per_token_loss else None,
             pg_collection=pg_collection,
+            force_all_reduce=force_all_reduce,
         )
 
     if not forward_only and config.fine_grained_activation_offloading:
@@ -1974,6 +1982,7 @@ def forward_backward_pipelining_without_interleaving(
     adjust_tensor_shapes_fn: Optional[Callable] = None,
     p2p_communicator: Optional[P2PCommunicator] = None,
     pg_collection: Optional[ProcessGroupCollection] = None,
+    force_all_reduce: Optional[bool] = False,
 ):
     """Run non-interleaved 1F1B schedule, with communication between pipeline
     stages. Returns dictionary with losses if the last stage, empty dict otherwise."""
@@ -2301,6 +2310,7 @@ def enable_grad_sync():
             [model],
             total_num_tokens if config.calculate_per_token_loss else None,
             pg_collection=pg_collection,
+            force_all_reduce=force_all_reduce,
         )
 
     if not forward_only and config.fine_grained_activation_offloading:
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 9fadb3f9900..3c2a0e52f71 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -2597,6 +2597,10 @@ def _add_checkpointing_args(parser):
                        help='Output directory to save checkpoints to.')
     group.add_argument('--save-interval', '--persistent-save-interval', type=int, default=None,
                        help='Number of iterations between persistent checkpoint saves.')
+    group.add_argument('--save-wgrads-interval', type=int, default=None,
+                       help='Number of iterations between wgrad (main_grad) saves.')
+    group.add_argument('--save-dgrads-interval', type=int, default=None,
+                       help='Number of iterations between dgrad saves.')
     group.add_argument('--save-retain-interval', type=int, default=None,
                        help='Number of iterations between retained checkpoints (other'
                        'checkpoints _except the last checkpoint_ are automatically deleted).')
diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py
index a4605121f82..6a0cc570da3 100644
--- a/megatron/training/checkpointing.py
+++ b/megatron/training/checkpointing.py
@@ -432,6 +432,40 @@ def _build_sharded_state_dict_metadata(args: Namespace, dp_cp_group: Optional[to
     metadata['dp_cp_group'] = dp_cp_group
     return metadata
 
+
+def save_grads(state_dict, iteration, grad_label):
+    args = get_args()
+
+    # Persist state_dict of grads onto disk. In case of wgrads, this collection should be
+    # performed before the grads are cleared but after they are reduced.
+    # NOTE: Non-expert layers will be duplicated, but this can be handled in postprocessing.
+
+    print_rank_0(f"  [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] saving {grad_label} "
+                 f"from iteration {iteration:7d}")
+
+    if mpu.get_expert_data_parallel_rank() == 0:
+        # Create saving directory.
+        ep_rank = mpu.get_expert_model_parallel_rank()
+        pp_rank = mpu.get_pipeline_model_parallel_rank()
+        tp_rank = mpu.get_tensor_model_parallel_rank()
+        assert args.save is not None
+        assert iteration is not None
+        save_dir = os.path.join(args.save, grad_label, f"iter_{iteration:07d}")
+        os.makedirs(save_dir, exist_ok=True)
+
+        # Save state_dict.
+        checkpoint_name = f"mp_rank_{tp_rank:02d}"
+        if mpu.get_pipeline_model_parallel_world_size() > 1:
+            checkpoint_name += f"_{pp_rank:03d}"
+        if mpu.get_expert_model_parallel_world_size() > 1:
+            checkpoint_name += f"_{ep_rank:03d}"
+        full_save_path = os.path.join(save_dir, f"{checkpoint_name}.pth")
+        torch.save(state_dict, full_save_path)
+
+    print_rank_0(f"  [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] saved {grad_label} "
+                 f"from iteration {iteration:7d}")
+
+
 def save_checkpoint(iteration, model, optimizer, opt_param_scheduler, num_floating_point_operations_so_far,
                     checkpointing_context=None, pipeline_rank=None, expert_rank=None, tensor_rank=None, pipeline_parallel=None, expert_parallel=None, non_persistent_ckpt=False,
                     train_data_iterator=None, preprocess_common_state_dict_fn = None, release=False, tp_group: Optional[torch.distributed.ProcessGroup] = None, pp_group: Optional[torch.distributed.ProcessGroup] = None, dp_cp_group: Optional[torch.distributed.ProcessGroup] = None):
diff --git a/megatron/training/dgrad_logging.py b/megatron/training/dgrad_logging.py
new file mode 100644
index 00000000000..2cd876ba848
--- /dev/null
+++ b/megatron/training/dgrad_logging.py
@@ -0,0 +1,114 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+
+"""dgrad logging using backward hooks."""
+
+import torch
+import torch.nn as nn
+
+from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
+
+from .checkpointing import save_grads
+from .utils import unwrap_model
+
+
+def _get_linear_types():
+    """Build tuple of linear layer types to capture gradients from."""
+    types = [nn.Linear, nn.Embedding, ColumnParallelLinear, RowParallelLinear]
+
+    # Add Transformer Engine layers if available.
+    try:
+        from megatron.core.extensions.transformer_engine import (
+            TELinear,
+            TEColumnParallelLinear,
+            TERowParallelLinear,
+            TELayerNormColumnParallelLinear,
+        )
+        types.extend([TELinear, TEColumnParallelLinear, TERowParallelLinear,
+                      TELayerNormColumnParallelLinear])
+    except ImportError:
+        pass
+
+    try:
+        from megatron.core.extensions.transformer_engine import (
+            TEGroupedLinear,
+            TEColumnParallelGroupedLinear,
+            TERowParallelGroupedLinear,
+        )
+        if TEGroupedLinear is not None:
+            types.extend([TEGroupedLinear, TEColumnParallelGroupedLinear,
+                          TERowParallelGroupedLinear])
+    except ImportError:
+        pass
+
+    return tuple(types)
+
+
+LINEAR_TYPES = _get_linear_types()
+
+
+class DataGradLogger:
+    """Captures and saves gradients from all linear layers using backward hooks."""
+
+    def __init__(self):
+        self._dgrads_state_dict = {}
+        self._hooks = []
+
+    def _make_hook(self, name: str):
+        """Create a backward hook for a named module."""
+        def hook(_, grad_input, grad_output):
+            for idx, grad in enumerate(grad_output):
+                if grad is not None:
+                    self._dgrads_state_dict[f"{name}/output{idx}"] = grad.detach().cpu()
+            for idx, grad in enumerate(grad_input):
+                if grad is not None:
+                    self._dgrads_state_dict[f"{name}/input{idx}"] = grad.detach().cpu()
+        return hook
+
+    def save(self, iteration: int):
+        """Save captured gradients to disk and clear the buffer."""
+        if not self._dgrads_state_dict:
+            return
+        save_grads(self._dgrads_state_dict, iteration, "dgrads")
+        self._dgrads_state_dict.clear()
+
+    def register_hooks(self, model: torch.nn.Module):
+        """Find and register hooks on all linear layers."""
+        assert len(self._hooks) == 0
+        for model_chunk_id, model_chunk in enumerate(model):
+            unwrapped_model_chunk = unwrap_model(model_chunk)
+            for name, module in unwrapped_model_chunk.named_modules():
+                if isinstance(module, LINEAR_TYPES):
+                    name = f"model_chunk{model_chunk_id}/{name}"
+                    handle = module.register_full_backward_hook(self._make_hook(name))
+                    self._hooks.append(handle)
+
+    def remove_hooks(self):
+        """Remove all registered hooks."""
+        for handle in self._hooks:
+            handle.remove()
+        self._hooks.clear()
+
+
+_LOGGER = None
+
+
+def enable_dgrad_logging(model: torch.nn.Module):
+    """Enable dgrad logging on a model."""
+    global _LOGGER
+    if _LOGGER is None:
+        _LOGGER = DataGradLogger()
+    _LOGGER.register_hooks(model)
+
+
+def disable_dgrad_logging():
+    """Disable dgrad logging on a model."""
+    global _LOGGER
+    assert _LOGGER is not None
+    _LOGGER.remove_hooks()
+
+
+def save_dgrads(iteration: int):
+    """Save dgrads to disk."""
+    global _LOGGER
+    assert _LOGGER is not None
+    _LOGGER.save(iteration)
diff --git a/megatron/training/training.py b/megatron/training/training.py
index 08a655ba80d..9e05fc6bf9c 100644
--- a/megatron/training/training.py
+++ b/megatron/training/training.py
@@ -100,7 +100,7 @@ def set_startup_timestamps(program_start=None, main_entry=None):
 )
 from megatron.core.optimizer import get_standard_config_overrides
 from megatron.training.checkpointing import load_checkpoint
-from megatron.training.checkpointing import save_checkpoint
+from megatron.training.checkpointing import save_checkpoint, save_grads
 from megatron.training.checkpointing import checkpoint_exists
 from megatron.training.checkpointing import get_loaded_iteration
 from megatron.core.full_cuda_graph import FullCudaGraphWrapper
@@ -187,6 +187,7 @@ def set_startup_timestamps(program_start=None, main_entry=None):
     get_energy_monitor,
 )
 from . import one_logger_utils
+from .dgrad_logging import enable_dgrad_logging, disable_dgrad_logging, save_dgrads
 
 from . import ft_integration
 
@@ -1604,16 +1605,23 @@ def dummy_train_step(data_iterator):
             batch = get_batch_on_this_cp_rank(batch)
 
 
-def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_scheduler, config, forward_backward_func):
+def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_scheduler, config, forward_backward_func, iteration=None):
     """Single training step."""
     args = get_args()
     timers = get_timers()
 
     rerun_state_machine = get_rerun_state_machine()
+    save_dgrads_in_this_iteration = (args.save_dgrads_interval is not None and
+                                     (iteration + 1) % args.save_dgrads_interval == 0)
+    # If saving main_grads in this iteration, then all-reduce instead of reduce-scatter.
+    save_wgrads_in_this_iteration = (args.save_wgrads_interval is not None and
+                                     (iteration + 1) % args.save_wgrads_interval == 0)
     while rerun_state_machine.should_run_forward_backward(data_iterator):
         # Set grad to zero.
         for model_chunk in model:
             model_chunk.zero_grad_buffer()
+            # Set force_all_reduce field for this iteration.
+            model_chunk.force_all_reduce = save_wgrads_in_this_iteration
         optimizer.zero_grad()
 
         if has_nvidia_modelopt:
@@ -1636,6 +1644,8 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch
                     optim_instance._copy_main_params_to_param_buffer()
 
         # Forward pass.
+        if save_dgrads_in_this_iteration:
+            enable_dgrad_logging(model)
         losses_reduced = forward_backward_func(
             forward_step_func=forward_step_func,
             data_iterator=data_iterator,
@@ -1646,7 +1656,32 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch
             decoder_seq_length=args.decoder_seq_length,
             forward_only=False,
             adjust_tensor_shapes_fn=adjust_tensor_shapes_fn,
+            force_all_reduce=save_wgrads_in_this_iteration,
         )
+        if save_dgrads_in_this_iteration:
+            save_dgrads(iteration + 1)
+            disable_dgrad_logging()
+
+        # Reset force_all_reduce field.
+        for model_chunk in model:
+            model_chunk.force_all_reduce = False
+
+    # Checkpoint main_grads.
+    if save_wgrads_in_this_iteration:
+        # Collect state_dict of wgrads (each param's .main_grad field).
+        state_dict = {}
+        for model_chunk_id, model_chunk in enumerate(model):
+            model_chunk_name = f"model_chunk{model_chunk_id}"
+            state_dict[model_chunk_name] = {}
+            unwrapped_model_chunk = unwrap_model(model_chunk)
+            for param_name, param in unwrapped_model_chunk.named_parameters():
+                if getattr(param, "main_grad", None) is not None:
+                    main_grad_on_cpu = param.main_grad.cpu()
+                    state_dict[model_chunk_name][param_name] = main_grad_on_cpu
+
+        # iteration is 0-indexed, move to 1-indexed for checkpoint name and logging.
+        save_grads(state_dict, iteration + 1, "wgrads")
+
     should_checkpoint, should_exit, exit_code = rerun_state_machine.should_checkpoint_and_exit()
     if should_exit:
         return {}, True, should_checkpoint, should_exit, exit_code, None, None, 0
@@ -2748,7 +2783,7 @@ def get_e2e_base_metrics():
             num_zeros_in_grad,
             max_attention_logit,
         ) = train_step(
-            forward_step_func, train_data_iterator, model, optimizer, opt_param_scheduler, config, forward_backward_func
+            forward_step_func, train_data_iterator, model, optimizer, opt_param_scheduler, config, forward_backward_func, iteration=iteration
         )
         ft_integration.on_training_step_end()
         if should_checkpoint:

From 8419801c7917e62fb74040ffedc6d34bc53b4240 Mon Sep 17 00:00:00 2001
From: Deepak Narayanan <dnarayanan@nvidia.com>
Date: Thu, 22 Jan 2026 11:04:58 -0800
Subject: [PATCH 2/5] Address comments

---
 megatron/core/pipeline_parallel/schedules.py | 6 ++++--
 megatron/training/training.py                | 3 +--
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/megatron/core/pipeline_parallel/schedules.py b/megatron/core/pipeline_parallel/schedules.py
index 07b8c4c54e9..edca62be375 100644
--- a/megatron/core/pipeline_parallel/schedules.py
+++ b/megatron/core/pipeline_parallel/schedules.py
@@ -125,8 +125,10 @@ def forward_step(data_iterator, model):
         respective list of shapes. Thus it is not used in the other forward-backward functions
         which have different shape handling.
 
-    force_all_reduce (bool, optional): Run all-reduce for this iteration instead of maybe
-        reduce-scatter to ensure all ranks have fully reduced gradients.
+    force_all_reduce (bool, optional): If true, force use of all-reduce for gradient reduction
+        instead of reduce-scatter (if using distributed optimizer) in this iteration to ensure all
+        data-parallel ranks have fully reduced gradients. This is useful for easier wgrad saving
+        (can just inspect DP replica 0 to get full set of wgrads for entire model).
 
     Args:
         pp_size (Optional[int]): Pipeline model parallel size to use.
diff --git a/megatron/training/training.py b/megatron/training/training.py
index 9e05fc6bf9c..7f5815106a3 100644
--- a/megatron/training/training.py
+++ b/megatron/training/training.py
@@ -1613,14 +1613,13 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch
     rerun_state_machine = get_rerun_state_machine()
     save_dgrads_in_this_iteration = (args.save_dgrads_interval is not None and
                                      (iteration + 1) % args.save_dgrads_interval == 0)
-    # If saving main_grads in this iteration, then all-reduce instead of reduce-scatter.
     save_wgrads_in_this_iteration = (args.save_wgrads_interval is not None and
                                      (iteration + 1) % args.save_wgrads_interval == 0)
     while rerun_state_machine.should_run_forward_backward(data_iterator):
         # Set grad to zero.
         for model_chunk in model:
             model_chunk.zero_grad_buffer()
-            # Set force_all_reduce field for this iteration.
+            # If saving main_grads in this iteration, then all-reduce instead of reduce-scatter.
             model_chunk.force_all_reduce = save_wgrads_in_this_iteration
         optimizer.zero_grad()
 

From ab388060daa0b18fa2d9a56be27a4039f1111a2a Mon Sep 17 00:00:00 2001
From: Deepak Narayanan <dnarayanan@nvidia.com>
Date: Thu, 22 Jan 2026 11:19:19 -0800
Subject: [PATCH 3/5] Make dgrad and wgrad save format identical

Signed-off-by: Deepak Narayanan <dnarayanan@nvidia.com>
---
 megatron/training/checkpointing.py |  4 +++-
 megatron/training/dgrad_logging.py | 19 ++++++++++++-------
 megatron/training/training.py      |  4 ++--
 3 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py
index 6a0cc570da3..138cffd4ec5 100644
--- a/megatron/training/checkpointing.py
+++ b/megatron/training/checkpointing.py
@@ -10,6 +10,7 @@
 import threading
 import types
 from argparse import Namespace
+from datetime import datetime
 from enum import Enum, auto
 from logging import getLogger
 from pathlib import Path
@@ -460,7 +461,8 @@ def save_grads(state_dict, iteration, grad_label):
         if mpu.get_expert_model_parallel_world_size() > 1:
             checkpoint_name += f"_{ep_rank:03d}"
         full_save_path = os.path.join(save_dir, f"{checkpoint_name}.pth")
-        torch.save(state_dict, full_save_path)
+        # Convert back to dict (e.g., from collections.defaultdict) for easy loading later.
+        torch.save(dict(state_dict), full_save_path)
 
     print_rank_0(f"  [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] saved {grad_label} "
                  f"from iteration {iteration:7d}")
diff --git a/megatron/training/dgrad_logging.py b/megatron/training/dgrad_logging.py
index 2cd876ba848..407d5869587 100644
--- a/megatron/training/dgrad_logging.py
+++ b/megatron/training/dgrad_logging.py
@@ -2,6 +2,7 @@
 
 """dgrad logging using backward hooks."""
 
+from collections import defaultdict
 import torch
 import torch.nn as nn
 
@@ -50,18 +51,20 @@ class DataGradLogger:
     """Captures and saves gradients from all linear layers using backward hooks."""
 
     def __init__(self):
-        self._dgrads_state_dict = {}
+        self._dgrads_state_dict = defaultdict(dict)
         self._hooks = []
 
-    def _make_hook(self, name: str):
+    def _make_hook(self, model_chunk_name: str, module_name: str):
         """Create a backward hook for a named module."""
         def hook(_, grad_input, grad_output):
             for idx, grad in enumerate(grad_output):
                 if grad is not None:
-                    self._dgrads_state_dict[f"{name}/output{idx}"] = grad.detach().cpu()
+                    grad_name = f"{module_name}/output{idx}"
+                    self._dgrads_state_dict[model_chunk_name][grad_name] = grad.detach().cpu()
             for idx, grad in enumerate(grad_input):
                 if grad is not None:
-                    self._dgrads_state_dict[f"{name}/input{idx}"] = grad.detach().cpu()
+                    grad_name = f"{module_name}/input{idx}"
+                    self._dgrads_state_dict[model_chunk_name][grad_name] = grad.detach().cpu()
         return hook
 
     def save(self, iteration: int):
@@ -76,10 +79,12 @@ def register_hooks(self, model: torch.nn.Module):
         assert len(self._hooks) == 0
         for model_chunk_id, model_chunk in enumerate(model):
             unwrapped_model_chunk = unwrap_model(model_chunk)
-            for name, module in unwrapped_model_chunk.named_modules():
+            for module_name, module in unwrapped_model_chunk.named_modules():
                 if isinstance(module, LINEAR_TYPES):
-                    name = f"model_chunk{model_chunk_id}/{name}"
-                    handle = module.register_full_backward_hook(self._make_hook(name))
+                    model_chunk_name = f"model_chunk{model_chunk_id}"
+                    handle = module.register_full_backward_hook(
+                        self._make_hook(model_chunk_name, module_name)
+                    )
                     self._hooks.append(handle)
 
     def remove_hooks(self):
diff --git a/megatron/training/training.py b/megatron/training/training.py
index 7f5815106a3..7ec9e7f6e8c 100644
--- a/megatron/training/training.py
+++ b/megatron/training/training.py
@@ -31,6 +31,7 @@ def set_startup_timestamps(program_start=None, main_entry=None):
         _STARTUP_TIMESTAMPS['main_entry'] = main_entry
 
 
+from collections import defaultdict
 import copy
 import dataclasses
 from datetime import datetime, timedelta
@@ -1668,10 +1669,9 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch
     # Checkpoint main_grads.
     if save_wgrads_in_this_iteration:
         # Collect state_dict of wgrads (each param's .main_grad field).
-        state_dict = {}
+        state_dict = defaultdict(dict)
         for model_chunk_id, model_chunk in enumerate(model):
             model_chunk_name = f"model_chunk{model_chunk_id}"
-            state_dict[model_chunk_name] = {}
             unwrapped_model_chunk = unwrap_model(model_chunk)
             for param_name, param in unwrapped_model_chunk.named_parameters():
                 if getattr(param, "main_grad", None) is not None:

From b82bb60e52a049fc3ad2d505ff4037f79f7e0981 Mon Sep 17 00:00:00 2001
From: Deepak Narayanan <dnarayanan@nvidia.com>
Date: Thu, 22 Jan 2026 12:54:33 -0800
Subject: [PATCH 4/5] Address comments, including unit tests

Signed-off-by: Deepak Narayanan <dnarayanan@nvidia.com>
---
 megatron/training/checkpointing.py            | 14 ++---
 megatron/training/dgrad_logging.py            |  9 +--
 megatron/training/training.py                 |  4 +-
 .../distributed/test_param_and_grad_buffer.py | 55 +++++++++++++++++
 tests/unit_tests/test_training.py             | 59 +++++++++++++++++++
 5 files changed, 128 insertions(+), 13 deletions(-)

diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py
index 138cffd4ec5..93e8a8acd6a 100644
--- a/megatron/training/checkpointing.py
+++ b/megatron/training/checkpointing.py
@@ -434,12 +434,12 @@ def _build_sharded_state_dict_metadata(args: Namespace, dp_cp_group: Optional[to
     return metadata
 
 
-def save_grads(state_dict, iteration, grad_label):
-    args = get_args()
+def save_grads(save_dir, state_dict, iteration, grad_label):
+    """Persist state_dict of grads onto disk. In case of wgrads, this collection should
+    be performed before the grads are cleared but after they are reduced.
 
-    # Persist state_dict of grads onto disk. In case of wgrads, this collection should be
-    # performed before the grads are cleared but after they are reduced.
-    # NOTE: Non-expert layers will be duplicated, but this can be handled in postprocessing.
+    NOTE: wgrads for non-expert layers will be duplicated if using expert parallelism, but
+    this can be handled in postprocessing."""
 
     print_rank_0(f"  [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] saving {grad_label} "
                  f"from iteration {iteration:7d}")
@@ -449,9 +449,9 @@ def save_grads(state_dict, iteration, grad_label):
         ep_rank = mpu.get_expert_model_parallel_rank()
         pp_rank = mpu.get_pipeline_model_parallel_rank()
         tp_rank = mpu.get_tensor_model_parallel_rank()
-        assert args.save is not None
+        assert save_dir is not None
         assert iteration is not None
-        save_dir = os.path.join(args.save, grad_label, f"iter_{iteration:07d}")
+        save_dir = os.path.join(save_dir, grad_label, f"iter_{iteration:07d}")
         os.makedirs(save_dir, exist_ok=True)
 
         # Save state_dict.
diff --git a/megatron/training/dgrad_logging.py b/megatron/training/dgrad_logging.py
index 407d5869587..bcab4b9f859 100644
--- a/megatron/training/dgrad_logging.py
+++ b/megatron/training/dgrad_logging.py
@@ -50,7 +50,8 @@ def _get_linear_types():
 class DataGradLogger:
     """Captures and saves gradients from all linear layers using backward hooks."""
 
-    def __init__(self):
+    def __init__(self, save_dir: str):
+        self._save_dir = save_dir
         self._dgrads_state_dict = defaultdict(dict)
         self._hooks = []
 
@@ -71,7 +72,7 @@ def save(self, iteration: int):
         """Save captured gradients to disk and clear the buffer."""
         if not self._dgrads_state_dict:
             return
-        save_grads(self._dgrads_state_dict, iteration, "dgrads")
+        save_grads(self._save_dir, self._dgrads_state_dict, iteration, "dgrads")
         self._dgrads_state_dict.clear()
 
     def register_hooks(self, model: torch.nn.Module):
@@ -97,11 +98,11 @@ def remove_hooks(self):
 _LOGGER = None
 
 
-def enable_dgrad_logging(model: torch.nn.Module):
+def enable_dgrad_logging(model: torch.nn.Module, save_dir: str):
     """Enable dgrad logging on a model."""
     global _LOGGER
     if _LOGGER is None:
-        _LOGGER = DataGradLogger()
+        _LOGGER = DataGradLogger(save_dir)
     _LOGGER.register_hooks(model)
 
 
diff --git a/megatron/training/training.py b/megatron/training/training.py
index 7ec9e7f6e8c..b7040e7bbb9 100644
--- a/megatron/training/training.py
+++ b/megatron/training/training.py
@@ -1645,7 +1645,7 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch
 
         # Forward pass.
         if save_dgrads_in_this_iteration:
-            enable_dgrad_logging(model)
+            enable_dgrad_logging(model, args.save)
         losses_reduced = forward_backward_func(
             forward_step_func=forward_step_func,
             data_iterator=data_iterator,
@@ -1679,7 +1679,7 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch
                     state_dict[model_chunk_name][param_name] = main_grad_on_cpu
 
         # iteration is 0-indexed, move to 1-indexed for checkpoint name and logging.
-        save_grads(state_dict, iteration + 1, "wgrads")
+        save_grads(args.save, state_dict, iteration + 1, "wgrads")
 
     should_checkpoint, should_exit, exit_code = rerun_state_machine.should_checkpoint_and_exit()
     if should_exit:
diff --git a/tests/unit_tests/distributed/test_param_and_grad_buffer.py b/tests/unit_tests/distributed/test_param_and_grad_buffer.py
index ac0c6a6c422..609b2cc5a71 100644
--- a/tests/unit_tests/distributed/test_param_and_grad_buffer.py
+++ b/tests/unit_tests/distributed/test_param_and_grad_buffer.py
@@ -3,6 +3,7 @@
 import contextlib
 import math
 from typing import Optional
+from unittest import mock
 
 import pytest
 import torch
@@ -265,3 +266,57 @@ def test_grad_sync(
             param_and_grad_buffer.grad_data.data.fill_(1.0)
 
     Utils.destroy_model_parallel()
+
+
+@pytest.mark.parametrize("force_all_reduce", [False, True])
+def test_force_all_reduce_uses_correct_collective(force_all_reduce: bool):
+    """Test that force_all_reduce=True causes all-reduce to be used instead of reduce-scatter."""
+    Utils.initialize_model_parallel()
+
+    input_dim = 100
+    output_dim = 100
+    num_layers = 2
+    model, param_and_grad_buffer, _ = get_model_and_buffers(
+        input_dim=input_dim,
+        output_dim=output_dim,
+        num_layers=num_layers,
+        bias=True,
+        shared_embedding=False,
+        bucket_size=None,
+        use_distributed_optimizer=True,  # This normally uses reduce-scatter.
+        overlap_grad_reduce=False,
+        average_in_collective=False,
+    )
+
+    # Mock the collective operations to track which one is called.
+    with (
+        mock.patch('torch.distributed.all_reduce') as mock_all_reduce,
+        mock.patch(
+            'megatron.core.distributed.param_and_grad_buffer.dist_reduce_scatter_func'
+        ) as mock_reduce_scatter,
+    ):
+        # Set up the mocks to be no-ops.
+        mock_all_reduce.return_value = None
+        mock_reduce_scatter.return_value = None
+
+        # Trigger the grad sync via the DDP model's finish_grad_sync method.
+        model.finish_grad_sync(force_all_reduce=force_all_reduce)
+
+        if force_all_reduce:
+            # When force_all_reduce=True, all_reduce should be called.
+            assert (
+                mock_all_reduce.called
+            ), "Expected all_reduce to be called when force_all_reduce=True"
+            assert (
+                not mock_reduce_scatter.called
+            ), "Expected reduce_scatter NOT to be called when force_all_reduce=True"
+        else:
+            # When force_all_reduce=False with distributed optimizer, reduce_scatter should be called.
+            assert (
+                mock_reduce_scatter.called
+            ), "Expected reduce_scatter to be called when force_all_reduce=False"
+            assert (
+                not mock_all_reduce.called
+            ), "Expected all_reduce NOT to be called when force_all_reduce=False"
+
+    Utils.destroy_model_parallel()
diff --git a/tests/unit_tests/test_training.py b/tests/unit_tests/test_training.py
index fef4bfbc5ef..2fd85724826 100644
--- a/tests/unit_tests/test_training.py
+++ b/tests/unit_tests/test_training.py
@@ -1,10 +1,16 @@
 # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
+from collections import defaultdict
+from pathlib import Path
 from types import SimpleNamespace
 
+import torch
+
+from megatron.training.checkpointing import save_grads
 from megatron.training.global_vars import set_args
 from megatron.training.tokenizer.tokenizer import _vocab_size_with_padding
 from megatron.training.training import build_train_valid_test_data_iterators
+from tests.unit_tests.dist_checkpointing import TempNamedDir
 from tests.unit_tests.test_utilities import Utils
 
 
@@ -74,3 +80,56 @@ def old_round_impl(after, multiple):
 
     def teardown_method(self, method):
         Utils.destroy_model_parallel()
+
+
+class TestSaveGrads:
+    """Tests for the save_grads function."""
+
+    def setup_method(self, method):
+        Utils.initialize_model_parallel(1, 1)
+
+    def teardown_method(self, method):
+        Utils.destroy_model_parallel()
+
+    def test_save_grads(self, tmp_path_dist_ckpt):
+        """Test that save_grads creates the correct directory structure and saves
+        state_dict correctly.
+
+        With TP=1, PP=1 on 8 GPUs, we have 8 DP ranks. Only the rank with
+        expert_data_parallel_rank==0 should save. All ranks verify the result.
+        """
+        save_dir = str(tmp_path_dist_ckpt / "test_save_grads")
+
+        with TempNamedDir(save_dir, sync=True) as save_dir:
+            # Create a mock state_dict with gradients (use deterministic values for reproducibility).
+            state_dict = defaultdict(dict)
+            state_dict["model_chunk0"]["layer.weight"] = torch.arange(16).reshape(4, 4).float()
+            state_dict["model_chunk0"]["layer.bias"] = torch.arange(4).float()
+
+            iteration = 100
+            grad_label = "wgrads"
+
+            # All ranks call save_grads, but only expert_data_parallel_rank==0 actually saves.
+            save_grads(save_dir, dict(state_dict), iteration, grad_label)
+
+            # Synchronize before checking results since only rank 0 saves.
+            torch.distributed.barrier()
+
+            # All ranks verify the file was created by rank 0.
+            expected_dir = Path(save_dir) / grad_label / f"iter_{iteration:07d}"
+            assert expected_dir.exists(), f"Expected directory {expected_dir} to exist"
+
+            expected_file = expected_dir / "mp_rank_00.pth"
+            assert expected_file.exists(), f"Expected file {expected_file} to exist"
+
+            # Verify saved content.
+            loaded = torch.load(expected_file)
+            assert "model_chunk0" in loaded
+            assert "layer.weight" in loaded["model_chunk0"]
+            assert "layer.bias" in loaded["model_chunk0"]
+            assert torch.equal(
+                loaded["model_chunk0"]["layer.weight"], state_dict["model_chunk0"]["layer.weight"]
+            )
+            assert torch.equal(
+                loaded["model_chunk0"]["layer.bias"], state_dict["model_chunk0"]["layer.bias"]
+            )

From 1dfb8b523d1f9f0a9ade97e5846aa54d532a2a98 Mon Sep 17 00:00:00 2001
From: Deepak Narayanan <dnarayanan@nvidia.com>
Date: Fri, 23 Jan 2026 11:42:58 -0800
Subject: [PATCH 5/5] Small clarification comment on implementation

Signed-off-by: Deepak Narayanan <dnarayanan@nvidia.com>
---
 megatron/training/dgrad_logging.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/megatron/training/dgrad_logging.py b/megatron/training/dgrad_logging.py
index bcab4b9f859..c046b4709fb 100644
--- a/megatron/training/dgrad_logging.py
+++ b/megatron/training/dgrad_logging.py
@@ -48,7 +48,10 @@ def _get_linear_types():
 
 
 class DataGradLogger:
-    """Captures and saves gradients from all linear layers using backward hooks."""
+    """Captures and saves gradients from all linear layers using backward hooks.
+    
+    NOTE: Right now, we only save the dgrads for the last microbatch in a batch on DP replica 0.
+    The code below would need to be extended to save dgrads for all microbatches in a batch."""
 
     def __init__(self, save_dir: str):
         self._save_dir = save_dir