Optimize memory usage of partial CUDA graphs

megatron/core/distributed/distributed_data_parallel.py

@@ -459,7 +459,8 @@ def hook(*unused):
 
             if param in self.param_to_bucket_group:
                 assert param.requires_grad
-                if self.ddp_config.overlap_grad_reduce:
+                cudagraph_wgrad_ready_event = getattr(param, '_cudagraph_wgrad_ready_event', None)
+                if self.ddp_config.overlap_grad_reduce and cudagraph_wgrad_ready_event is None:
                     assert (
                         param.grad is not None
                     ), 'param.grad being None is not safe when overlap_grad_reduce is True'

megatron/core/distributed/param_and_grad_buffer.py

@@ -586,6 +586,18 @@ def start_grad_sync(self, force_all_reduce: Optional[bool] = False):
             self.grad_reduce_handle is None
         ), "Should not have multiple communication calls outstanding at once"
 
+        # Local CUDA graph replay is asynchronous with respect to the outer
+        # autograd hooks. Wait before reading, scaling, or reducing gradients
+        # accumulated by a replay into this bucket.
+        current_stream = torch.cuda.current_stream()
+        waited_event_ids = set()
+        for bucket in self.buckets:
+            for param in bucket.params_list:
+                event = getattr(param, "_cudagraph_wgrad_ready_event", None)
+                if event is not None and id(event) not in waited_event_ids:
+                    current_stream.wait_event(event)
+                    waited_event_ids.add(id(event))
+
         # Copy accumulated .main_grad into communication buffer before collective if
         # .main_grad is not in .grad_data already (e.g., because we want to do local
         # gradient accumulation in a higher precision).

megatron/core/models/hybrid/hybrid_block.py

@@ -25,6 +25,7 @@
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.recompute import checkpointed_forward
 from megatron.core.transformer import TransformerConfig
+from megatron.core.transformer.cuda_graphs import annotate_first_last_layer
 from megatron.core.transformer.identity_op import IdentityOp
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.spec_utils import ModuleSpec, build_module
@@ -187,6 +188,9 @@ def __init__(
                     raise ValueError("unexpected layer_type")
             self.layers.append(layer)
 
+        if self.config.cuda_graph_impl == "local":
+            annotate_first_last_layer(self.layers)
+
         # Required for activation recomputation
         self.num_layers_per_pipeline_rank = len(self.layers)