fix: ensure GTP params are ready before CUDA graph replay

megatron/core/tensor_parallel/generalized_tensor_parallelism.py

@@ -40,6 +40,7 @@
     allocate_graph_wgrad_rings,
     cuda_graph_pool_allocation,
     register_capture_comm,
+    register_capture_params_to_ensure_ready,
     register_capture_wgrad_ring_slot,
 )
 from megatron.core.utils import ensure_params_ready, log_single_rank
@@ -1297,6 +1298,7 @@ def _all_gather_weight(
         #    Forward only: backward re-reads what forward published, and recompute runs inside
         #    backward where a dispatch could gather into the grad-aliased buffer.
         if fwd and not in_activation_recompute_phase():
+            register_capture_params_to_ensure_ready(weights)
             ensure_params_ready(weights)
 
         # 1. Transition state for async gathers. Skip during recompute-forward: it gathers

megatron/core/tensor_parallel/gtp_cuda_graphs.py

@@ -36,13 +36,15 @@ class GraphWgradRingSlot:
 
 @dataclass
 class GTPCaptureCommState:
-    """Asynchronous GTP work issued while capturing one CUDA graph."""
+    """GTP work and parameter-readiness dependencies issued while capturing one CUDA graph."""
 
     params: list = field(default_factory=list)
+    params_to_ensure_ready: list = field(default_factory=list)
     ag_streams: list = field(default_factory=list)
     rs_streams: list = field(default_factory=list)
     wgrad_ring_slots: list = field(default_factory=list)
     _param_ids: set = field(default_factory=set)
+    _param_ids_to_ensure_ready: set = field(default_factory=set)
     _ag_stream_ids: set = field(default_factory=set)
     _rs_stream_ids: set = field(default_factory=set)
     _wgrad_ring_slot_params: dict = field(default_factory=dict)
@@ -61,6 +63,14 @@ def register_comm(self, param, stream: torch.cuda.Stream, *, reduce_scatter: boo
             stream_ids.add(stream_id)
             streams.append(stream)
 
+    def register_params_to_ensure_ready(self, params: Iterable) -> None:
+        """Record parameters that must be published before this graph replays."""
+        for param in params:
+            param_id = id(param)
+            if param_id not in self._param_ids_to_ensure_ready:
+                self._param_ids_to_ensure_ready.add(param_id)
+                self.params_to_ensure_ready.append(param)
+
     def register_wgrad_ring_slot(self, slot: GraphWgradRingSlot, param) -> None:
         """Track slots used by this graph and reject unsafe intra-graph aliasing."""
         slot_id = id(slot)
@@ -85,6 +95,12 @@ def register_capture_comm(param, stream: torch.cuda.Stream, *, reduce_scatter: b
         _ACTIVE_CAPTURE_COMM_STATE.register_comm(param, stream, reduce_scatter=reduce_scatter)
 
 
+def register_capture_params_to_ensure_ready(params: Iterable) -> None:
+    """Record GTP parameter reads that need publication before graph replay."""
+    if _ACTIVE_CAPTURE_COMM_STATE is not None:
+        _ACTIVE_CAPTURE_COMM_STATE.register_params_to_ensure_ready(params)
+
+
 def register_capture_wgrad_ring_slot(slot: GraphWgradRingSlot, param) -> None:
     """Register a ring slot with the active capture, if one exists."""
     if _ACTIVE_CAPTURE_COMM_STATE is not None:

megatron/core/transformer/cuda_graphs.py

@@ -32,6 +32,7 @@
 from megatron.core.transformer.module import GraphableMegatronModule, MegatronModule
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.utils import (
+    ensure_params_ready,
     get_attr_wrapped_model,
     get_torch_version,
     is_te_min_version,
@@ -710,6 +711,7 @@ def delete_cuda_graphs():
         runner.fwd_graph = None
         runner.bwd_graph = None
         runner.mempool = None
+        runner._gtp_fwd_params_to_ensure_ready = ()
 
     # Reset global tracking state
     _CudagraphGlobalRecord.cudagraph_created = False
@@ -948,6 +950,8 @@ def __init__(
         # Persistent wgrad slots written by this graph. Replay waits for each slot's previous RS
         # reader before launching the graph.
         self._gtp_wgrad_ring_slots = []
+        # GTP weights read by this forward graph before their owning module pre-hooks execute.
+        self._gtp_fwd_params_to_ensure_ready = ()
 
         self.grad_enabled = need_backward and torch.is_grad_enabled()
         self.func = super(MegatronModule, self.base_module).__call__ if func is None else func
@@ -1335,8 +1339,14 @@ def clone_ten(ten):
                 if FREEZE_GC:
                     gc.freeze()
 
-                with torch.cuda.graph(
-                    self.fwd_graph, pool=self.mempool, capture_error_mode="thread_local"
+                capture_comm_context = (
+                    track_gtp_capture_comms() if self.gtp_remat else nullcontext(None)
+                )
+                with (
+                    capture_comm_context as capture_comms,
+                    torch.cuda.graph(
+                        self.fwd_graph, pool=self.mempool, capture_error_mode="thread_local"
+                    ),
                 ):
 
                     self._sync_against_side_streams(self.fwd_side_streams)
@@ -1355,6 +1365,11 @@ def clone_ten(ten):
                     if self.use_stream:
                         self.fwd_completion_event.record()
 
+                if self.gtp_remat:
+                    self._gtp_fwd_params_to_ensure_ready = tuple(
+                        capture_comms.params_to_ensure_ready
+                    )
+
                 # Unfreeze GC.
                 if FREEZE_GC:
                     gc.unfreeze()
@@ -1665,6 +1680,9 @@ def replay_graph_capture(self, is_first_microbatch, args, kwargs):
             error_msg = "CUDA graph argument mismatch:\n" + "\n".join(mismatch_errors)
             raise AssertionError(error_msg)
 
+        if self._gtp_fwd_params_to_ensure_ready:
+            ensure_params_ready(self._gtp_fwd_params_to_ensure_ready)
+
         inp_tensors = self.get_tensors(args, kwargs, check_types=False)
         if self.grad_enabled:
             func_args = inp_tensors + self.params_to_backprop

tests/unit_tests/distributed/test_param_readiness.py

@@ -5,13 +5,15 @@
 The GPU test in ``generalized_tensor_parallel/`` runs one configuration: healthy DDP, pre-hooks
 enabled, ``align_param_gather=False``. These cover the branches it never exercises --
 ``align_param_gather=True``, pre-hooks removed mid-sequence, a collected DDP or bucket group,
-and out-of-order bucket dedup -- which would be contrived to set up on real GPUs.
+CUDA-graph capture deferral, and out-of-order bucket dedup -- which would be contrived to set up
+on real GPUs.
 """
 
 from types import SimpleNamespace
 
 import torch
 
+import megatron.core.distributed.distributed_data_parallel as ddp_module
 from megatron.core.distributed.distributed_data_parallel import (
     DistributedDataParallel,
     _BucketParamReadyCallback,
@@ -70,6 +72,24 @@ class TestBucketParamReadiness:
     immediately and make every callback early-return -- turning these into vacuous passes.
     """
 
+    def test_cuda_graph_capture_defers_publication_until_replay(self, monkeypatch):
+        ddp = _fake_ddp()
+        bucket_group = _FakeBucketGroup()
+        ready_callback = _BucketParamReadyCallback(ddp, bucket_group)
+        param = torch.nn.Parameter(torch.zeros(1))
+        setattr(param, PARAM_READY_CALLBACK_ATTR, ready_callback)
+        capturing = [True]
+        monkeypatch.setattr(ddp_module, "is_graph_capturing", lambda: capturing[0])
+
+        ensure_params_ready([param])
+        assert bucket_group.finished == []
+        assert not bucket_group.param_gather_dispatched
+
+        capturing[0] = False
+        ensure_params_ready([param])
+        assert bucket_group.finished == [False]
+        assert bucket_group.param_gather_dispatched
+
     def test_publishes_once_then_is_idempotent(self):
         ddp = _fake_ddp()
         bucket_group = _FakeBucketGroup()

tests/unit_tests/generalized_tensor_parallel/test_gtp_basics.py

@@ -1591,3 +1591,40 @@ def forward(self, x):
             f"expected [original fwd=False, recompute fwd=True], got {observed}: "
             "in_activation_recompute_phase no longer tracks BF16 activation recompute"
         )
+
+
+class TestGTPCaptureParamReadiness:
+    def test_forward_gather_registers_params_before_ensuring_readiness(self, monkeypatch):
+        class StopAfterReadiness(Exception):
+            pass
+
+        first = object()
+        second = object()
+        param = type("Param", (), {"_debug_name": "weight", "_weights": (first, second)})()
+        readiness_calls = []
+
+        monkeypatch.setattr(gtp_module, "nvtx_range_push", lambda _: None)
+        monkeypatch.setattr(gtp_module, "in_activation_recompute_phase", lambda: False)
+
+        def stop_after_readiness(params):
+            readiness_calls.append(tuple(params))
+            raise StopAfterReadiness
+
+        monkeypatch.setattr(gtp_module, "ensure_params_ready", stop_after_readiness)
+
+        with gtp_cuda_graphs.track_gtp_capture_comms() as capture:
+            with pytest.raises(StopAfterReadiness):
+                GTPShardedParam._all_gather_weight(param, async_op=False, fwd=True)
+
+        assert readiness_calls == [(first, second)]
+        assert capture.params_to_ensure_ready == [first, second]
+
+    def test_param_readiness_records_unique_parameters(self):
+        first = object()
+        second = object()
+
+        with gtp_cuda_graphs.track_gtp_capture_comms() as capture:
+            gtp_cuda_graphs.register_capture_params_to_ensure_ready((first, second, first))
+            gtp_cuda_graphs.register_capture_params_to_ensure_ready((second,))
+
+        assert capture.params_to_ensure_ready == [first, second]

tests/unit_tests/transformer/test_cuda_graphs.py

@@ -8,6 +8,7 @@
 import torch
 from transformer_engine.pytorch.fp8 import check_fp8_support
 
+import megatron.core.transformer.cuda_graphs as cuda_graphs_module
 from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig
 from megatron.core.enums import ModelType
 from megatron.core.models.gpt.gpt_layer_specs import (
@@ -37,6 +38,8 @@
     CudaGraphManager,
     TECudaGraphHelper,
     _CudagraphGlobalRecord,
+    _CudagraphReplayNode,
+    _CudaGraphRunner,
     create_cudagraphs,
     delete_cuda_graphs,
 )
@@ -391,6 +394,32 @@ def test_deprecated_cuda_graph_scope_full_iteration_inference_noops_without_loca
         assert cfg.cuda_graph_scope is None
 
 
+class TestCudaGraphReplay:
+    def test_gtp_forward_ensures_captured_params_ready_before_replay(self, monkeypatch):
+        calls = []
+        first = object()
+        second = object()
+        runner = object.__new__(_CudaGraphRunner)
+        runner._gtp_fwd_params_to_ensure_ready = (first, second)
+        runner.grad_enabled = False
+        runner.fwd_graph_outputs = (object(),)
+        runner.get_mismatch_errors = lambda args, kwargs: []
+        runner.get_tensors = lambda args, kwargs, check_types: []
+        runner.to_list = lambda value: list(value) if isinstance(value, tuple) else [value]
+        monkeypatch.setattr(
+            cuda_graphs_module,
+            "ensure_params_ready",
+            lambda params: calls.append(("ready", tuple(params))),
+        )
+        monkeypatch.setattr(
+            _CudagraphReplayNode, "apply", lambda *args: calls.append("replay") or (object(),)
+        )
+
+        runner.replay_graph_capture(False, (), {})
+
+        assert calls == [("ready", (first, second)), "replay"]
+
+
 class TestParallelTransformerBlockCudagraphs:
     def setup_method(self, method):
         # initialize parallel state