fix(gtp): preserve scratch lifetime during full-iteration capture

megatron/core/full_cuda_graph.py

@@ -10,6 +10,7 @@
 import torch
 
 from megatron.core.rerun_state_machine import get_rerun_state_machine
+from megatron.core.tensor_parallel.gtp_symmetric_memory import symmetric_wgrad_pool
 from megatron.core.tensor_parallel.random import (
     cudagraph_needs_generator_registration,
     get_all_rng_states,
@@ -321,6 +322,7 @@ def __call__(self, *args, **kwargs):
                     FullCudaGraphWrapper.cuda_graph[training_str].register_generator_state(state)
             torch.cuda.synchronize()
             capture_stream = get_shared_capture_stream()
+            symmetric_wgrad_pool.reset_completed_reuse_events()
             rerun_state_machine = get_rerun_state_machine()
             with (
                 torch.cuda.graph(
@@ -358,6 +360,7 @@ def next_iter(self, stage):
 
     def reset_cuda_graph(self, stage=None):
         """Reset CUDA graph."""
+        captured = any(graph is not None for graph in FullCudaGraphWrapper.cuda_graph.values())
         if stage is None or stage == 'training':
             if FullCudaGraphWrapper.cuda_graph['training'] is not None:
                 del FullCudaGraphWrapper.cuda_graph['training']
@@ -372,4 +375,6 @@ def reset_cuda_graph(self, stage=None):
             FullCudaGraphWrapper.result['validation'] = None
             FullCudaGraphWrapper.curr_iteration['validation'] = 0
             get_rerun_state_machine().clear_deferred_validations("validation")
+        if captured:
+            symmetric_wgrad_pool.reset_completed_reuse_events()
         gc.collect()

megatron/core/tensor_parallel/generalized_tensor_parallelism.py

@@ -268,6 +268,13 @@ def classify_gtp_chains(model) -> None:
             continue
         param.chain_id = target
 
+        if _FULL_ITERATION and not param.prefetch_initialized:
+            # All producers and consumers belong to one graph. Internal events expose
+            # their ordering to graph-pool reuse; partial graphs need external handoffs.
+            param.ag_event = torch.cuda.Event()
+            param._recompute_ag_event = torch.cuda.Event()
+            param.rs_event = torch.cuda.Event()
+
         # Bwd-prefetch opt-out: embedding weight needs no bwd AG (wgrad is a
         # scatter-add on sharded rows, input has no dgrad) — saves one collective.
         if "embedding" in name:
@@ -329,18 +336,33 @@ def _alloc_symmetric_wgrad_buffer(weight, dtype, device) -> torch.Tensor:
 
 
 def _wgrad_pool_get(shape: tuple, dtype: torch.dtype, device) -> torch.Tensor:
-    """Get a pool buffer or allocate fresh, tagged so _wgrad_pool_put accepts only
-    pool-owned buffers (other callers fall through to the caching allocator on release)."""
-    key = (shape, dtype)
-    pool = _wgrad_buf_pool.get(key)
+    """Allocate wgrad scratch: from the eager pool outside capture, from the graph pool inside.
+
+    The decision keys on the capture state of the current stream rather than on the chain's
+    classification, because that is what determines where the storage lives:
+
+    * Capturing (full-iteration or partial CUDA graph): ``torch.empty`` lands in the graph's
+      private pool and the captured kernels keep that address for every replay. The eager pool
+      is neither read (an eager buffer baked into a graph would later be overwritten by eager
+      reuse) nor written (the buffer is left untagged, so ``_wgrad_pool_put`` ignores it).
+    * Not capturing (eager execution, including the warmup passes before capture): pop a
+      cached buffer of this (shape, dtype) if there is one, waiting on its reuse event, else
+      allocate; tag it so ``_wgrad_pool_put`` may return it. Whether it does return is decided
+      at release by the chain: UNGRAPHED chains recycle it, GRAPHED chains drop it after
+      ``record_stream`` (see ``_release_wgrad_scratch``).
+    """
+    capturing = torch.cuda.is_current_stream_capturing()
+    # Graph region: never borrow eager storage.
+    pool = None if capturing else _wgrad_buf_pool.get((shape, dtype))
     if pool:
         buf = pool.pop()
         ready_event = getattr(buf, "_gtp_wgrad_reuse_event", None)
         if ready_event is not None:
             torch.cuda.current_stream(device=device).wait_event(ready_event)
     else:
         buf = torch.empty(shape, dtype=dtype, device=device, requires_grad=False)
-    buf._from_gtp_wgrad_pool = True
+    # Only eager allocations may be recycled; graph-pool storage never enters the eager cache.
+    buf._from_gtp_wgrad_pool = not capturing
     return buf
 
 
@@ -1998,11 +2020,12 @@ def get_wgrad_tensor(self, *, persistent: bool = False):
             # is still busy. Allocating a second raises this bucket's high-water mark for good --
             # the pool never shrinks, so every buffer in it is live at the peak. Drain our own RS
             # so the alloc below recycles that buffer, trading one overlap for the permanent
-            # bytes. Capture-guarded: the branch is host-side.
+            # bytes. Full-iteration capture needs the same drain as warmup to reuse
+            # that one buffer. Partial graphs leave this scheduling to their runner.
             if (
                 GTP_CONFIG.async_reduction
                 and self._wgrad_rs_handle is not None
-                and not torch.cuda.is_current_stream_capturing()
+                and (_FULL_ITERATION or not torch.cuda.is_current_stream_capturing())
                 and not symmetric_wgrad_pool.has_free(
                     self._unsharded_shape_padded, self.main_grad.dtype, self.group
                 )
@@ -2103,15 +2126,18 @@ def _wait_reduce_scatter(self, finalize_grad=False) -> bool:
                     for w in self._weights:
                         self._handle_megatron_grad_accum(w)
                     self._already_finalized = True
-        self._release_wgrad_scratch(ready_event=self.rs_event)
+        self._release_wgrad_scratch(stream=rs_stream, ready_event=self.rs_event)
         return waited
 
-    def _release_wgrad_scratch(self, attrs=("_wgrad_input_bufs", "_rs_a2a_bufs"), ready_event=None):
+    def _release_wgrad_scratch(
+        self, attrs=("_wgrad_input_bufs", "_rs_a2a_bufs"), stream=None, ready_event=None
+    ):
         """Release the buffers a finished RS was reading.
 
         Its wgrad inputs, and the fp32-accum all-to-all scratch (input to the deferred FP32 sum,
-        so only free once the handle has been waited on). UNGRAPHED buffers go back to the pool;
-        GRAPHED just drops Python refs (addresses must stay stable for CG).
+        so only release once the handle has been waited on). Plain UNGRAPHED buffers return
+        to the manual pool, graph-pool buffers drop their Python references, and symmetric
+        buffers return to their registered pool. Reuse waits for the last reader.
         Work.wait() orders only its calling stream: pooled buffers carry the completion event
         so their next writer waits at checkout, without blocking unrelated compute here.
         """
@@ -2122,8 +2148,13 @@ def _release_wgrad_scratch(self, attrs=("_wgrad_input_bufs", "_rs_a2a_bufs"), re
             if not _chain_is_graphed(self.chain_id):
                 for buf in bufs:
                     _wgrad_pool_put(buf, ready_event=ready_event)
+            elif stream is not None:
+                # The FP32 sum is enqueued by handle.wait(). Keep its input alive until
+                # that side-stream read completes, including during graph capture.
+                for buf in bufs:
+                    buf.record_stream(stream)
             for buf in bufs:
-                # Return symm pool buffers (tag-gated no-op for plain and ring buffers).
+                # Tag-gated no-op for buffers outside the registered pool.
                 symmetric_wgrad_pool.free(buf, ready_event=ready_event)
             setattr(self, attr, None)
 
@@ -2771,12 +2802,14 @@ def wait_async_comms(
 
     This drains AG and recompute-AG handles and, unless ``skip_rs`` is set, wgrad RS handles.
     During CUDA graph capture, draining a handle materializes its wait in the producer graph and
-    records the corresponding external event. Because the event was constructed with
-    ``external=True``, capture emits event-record and event-wait nodes that preserve the dependency
+    records the corresponding event. Partial graphs use ``external=True`` to emit
+    event-record and event-wait nodes that preserve the dependency
     between separately captured graphs at replay. The host must launch the producer graph before
     the consumer graph so its record is the one observed by the wait. An AG handoff marker prevents
     the consumer from draining the same Work handle again; the consumer still waits on the event.
     When requested, an RS result is also accumulated into ``main_grad`` before graph completion.
+    Full-iteration graphs use internal events because the producer and consumer are in one
+    graph, exposing their dependencies to capture-time allocator reuse.
 
     Args:
         chain_id: If specified, only drain params on this chain.

megatron/core/tensor_parallel/gtp_symmetric_memory.py

@@ -231,6 +231,26 @@ def free(self, buf: torch.Tensor, ready_event: torch.cuda.Event | None = None) -
         flat._gtp_wgrad_reuse_event = ready_event
         self._free[(buf.numel(), buf.dtype, group.group_name)].append(flat)
 
+    def reset_completed_reuse_events(self) -> None:
+        """Drop every reuse fence, retaining the registered storage.
+
+        Call outside capture, when no eager reader can be in flight (the device synchronized):
+
+        * before a full-iteration capture: fences recorded by eager work cannot be waited on
+          from inside a capture, even once that work has completed. Fences left by an earlier
+          capture may be present too; they must not be queried, which is why this drops
+          unconditionally rather than checking completion;
+        * when the graphs are reset: training resumes eagerly, and an eager ``wait_event`` on a
+          fence recorded inside a capture fails with ``cudaErrorInvalidValue``.
+
+        Releases during capture still attach their normal fences, which become graph edges.
+        """
+        if torch.cuda.is_current_stream_capturing():
+            raise RuntimeError("Reset GTP symmetric reuse events outside CUDA-graph capture")
+        for bucket in self._free.values():
+            for buf in bucket:
+                buf._gtp_wgrad_reuse_event = None
+
     def clear(self) -> None:
         """Drop every cached buffer. Called at teardown, before the pools they alias go away."""
         self._free.clear()

tests/unit_tests/generalized_tensor_parallel/test_gtp_full_iteration_buffers.py

@@ -0,0 +1,115 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+import pytest
+import torch
+
+from megatron.core.full_cuda_graph import FullCudaGraphWrapper
+from megatron.core.tensor_parallel import generalized_tensor_parallelism as gtp
+from megatron.core.tensor_parallel.gtp_api import HAVE_GTP
+from tests.unit_tests.generalized_tensor_parallel.gtp_test_utils import (
+    _make_gtp_linear,
+    reset_gtp_globals,
+)
+from tests.unit_tests.test_utilities import Utils
+
+
+def test_capture_does_not_borrow_eager_wgrad_storage(monkeypatch):
+    """A graph must not keep a scratch address that eager work can later overwrite."""
+    Utils.initialize_distributed()
+    monkeypatch.setattr(gtp, "_wgrad_buf_pool", {})
+    eager = gtp._wgrad_pool_get((1024,), torch.float32, "cuda")
+    assert eager._from_gtp_wgrad_pool
+    gtp._wgrad_pool_put(eager)
+    result = torch.empty_like(eager)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        scratch = gtp._wgrad_pool_get(eager.shape, eager.dtype, eager.device)
+        scratch.fill_(3)
+        result.copy_(scratch)
+        # Graph-owned scratch is never cached for eager reuse, whatever release path sees it.
+        assert not scratch._from_gtp_wgrad_pool
+        gtp._wgrad_pool_put(scratch)
+    assert scratch.data_ptr() != eager.data_ptr()
+    assert gtp._wgrad_pool_get(eager.shape, eager.dtype, eager.device) is eager
+    assert not gtp._wgrad_buf_pool.get((eager.shape, eager.dtype))
+    for _ in range(3):
+        eager.fill_(9)
+        graph.replay()
+        torch.testing.assert_close(result, torch.full_like(result, 3))
+
+
+@pytest.mark.skipif(not HAVE_GTP, reason="GTP requires TransformerEngine >= 2.19")
+def test_full_iteration_fp32_accumulation_scratch_lifetime(monkeypatch):
+    """Captured send scratch must outlive its RS-stream reader under allocation pressure.
+
+    The RS stream still reads the compute-stream send buffer after _wait_reduce_scatter
+    releases it; without last-reader tracking the graph pool hands that storage to the next
+    compute-stream allocation. (The all-to-all scratch is allocated on the RS stream itself,
+    so compute-stream pressure cannot alias it.)
+    """
+    Utils.initialize_distributed()
+    group = torch.distributed.group.WORLD
+    world_size = group.size()
+    if world_size <= 2:
+        pytest.skip("FP32-accumulation reduce-scatter requires more than two ranks")
+    assert not gtp.is_gtp_symm_pool_registered(group)
+    monkeypatch.setattr(gtp, "_wgrad_buf_pool", {})
+    monkeypatch.setattr(gtp, "_FULL_ITERATION", True)
+    monkeypatch.setattr(gtp.GTP_CONFIG, "async_reduction", True)
+    monkeypatch.setattr(gtp.GTP_CONFIG, "reduce_scatter_with_fp32_accumulation", True)
+    monkeypatch.setattr(gtp.GTP_CONFIG, "calculate_per_token_loss", False)
+    monkeypatch.setattr("megatron.core.full_cuda_graph.get_all_rng_states", lambda: {})
+    layer = _make_gtp_linear(128, 32 * world_size, group)
+    gtp.classify_gtp_chains(layer)
+    weight = layer.weight
+    weight.main_grad = torch.zeros_like(weight)
+    source = torch.empty(weight._unsharded_shape, dtype=torch.bfloat16, device="cuda")
+    source.fill_(group.rank())
+    calls = []
+
+    def forward_backward_func(**_):
+        weight.main_grad.zero_()
+        buffer = weight.get_wgrad_tensor()
+        buffer.copy_(source)
+        weight.wgrad_reduce_scatter(buffer, async_op=True)
+        del buffer
+        # This list is populated only by the non-registered all-to-all/FP32-sum path.
+        assert len(weight._rs_a2a_bufs) == 1
+        calls.append(torch.cuda.is_current_stream_capturing())
+        rs_stream = gtp.get_rs_stream(weight.chain_id, group)
+        with torch.cuda.stream(rs_stream):
+            torch.cuda._sleep(5_000_000)
+        weight._wait_reduce_scatter(finalize_grad=True)
+        assert weight._rs_a2a_bufs is None and weight._wgrad_input_bufs is None
+        # Do not join the RS stream until after these overwrites. Incorrectly freed send
+        # scratch can be recycled while the reduce-scatter is still reading it.
+        pressure = [torch.empty_like(source) for _ in range(8)]
+        for allocation in pressure:
+            allocation.fill_(-100)
+        torch.cuda.current_stream().wait_stream(rs_stream)
+        return [{"loss": weight.main_grad}]
+
+    wrapper = FullCudaGraphWrapper(forward_backward_func, cuda_graph_warmup_steps=1)
+    wrapper.reset_cuda_graph()
+    kwargs = dict(
+        model=[torch.nn.Identity()],
+        data_iterator=None,
+        num_microbatches=1,
+        seq_length=1,
+        forward_only=False,
+    )
+    try:
+        wrapper(**kwargs)
+        wrapper(**kwargs)
+        for step in range(10):
+            source.fill_(group.rank() + step)
+            result = wrapper(**kwargs)
+            torch.cuda.synchronize()
+            expected = step + (world_size - 1) / 2
+            torch.testing.assert_close(
+                result[0]["loss"], torch.full_like(weight.main_grad, expected), atol=0, rtol=0
+            )
+        assert calls == [False, True], "Replay must not repeat Python scratch allocation"
+    finally:
+        torch.cuda.synchronize()
+        wrapper.reset_cuda_graph()