[fix] Give same-key GTP chain neighbours distinct gather buffers

docs/api-guide/core/generalized_tensor_parallel.md

@@ -342,7 +342,7 @@ The figure visualizes the per-class split from the list above: green = resolves
 
 Two distinct pools with explicit lifecycle rules:
 
-- **`GTPWeightCache`** (AG/RS output buffers) — ticket-based, keyed on `(shape, dtype, fwd, expert_idx, reduce_scatter)`. Same-shape buffers across layers are shared. Tickets persistent; buffer allocated lazily on first `get()`; addresses stable across iterations for CG replay.
+- **`GTPWeightCache`** (AG/RS output buffers) — ticket-based, keyed on `(shape, dtype, fwd, expert_idx, reduce_scatter)`. Same-shape buffers across layers are shared, **except between chain neighbours** — one-step-ahead keeps `prev_w` and the current weight live at once, so `_ensure_distinct_buffer_from_prev` folds a parity bit into the key when the two would collide, at the cost of one extra buffer for the second of the pair. Normally inert (neighbours are different weight roles, hence different shapes); it fires when CG capture leaves two same-shaped weights adjacent — embedding + output_layer alone in the `UNGRAPHED` chain. Tickets persistent; buffer allocated lazily on first `get()`; addresses stable across iterations for CG replay.
 - **`_wgrad_buf_pool`** (wgrad-GEMM output recycling) — holds the **full, unsharded** wgrad-GEMM output buffer (shape `_unsharded_shape`, dtype `main_grad.dtype` — fp32 when `grad_reduce_in_fp32`, else bf16). The TE backward writes the wgrad into it via `main_grad_func = weight.grad_buffer` (a `DistributedWeight` protocol method backed by `get_wgrad_tensor`; it is a *scratch*, distinct from the sharded `param.main_grad`); the protocol's `finalize_group_grads` (backed by `wgrad_reduce_scatter`) then reduce-scatters it down to the shard and the buffer is returned here. This is a full-weight-shaped fp32/bf16 transient — one of the larger per-weight buffers — and is **precision-independent** (wgrad is always computed in high precision), so it is identical in BF16 vs MXFP8 runs. Buffers are tagged `_from_gtp_wgrad_pool=True` at `_wgrad_pool_get`; `_wgrad_pool_put` no-ops on foreign buffers (fresh allocs from Megatron `layers.py` or aten F.embedding bwd) → caching allocator handles those, so the pool never accumulates untagged buffers.
 
 #### Overlap design summary
@@ -527,7 +527,8 @@ Three consequences:
   - the weight cache keys **one buffer per `(shape, dtype, expert_idx)`**, which assumes at most one same-key weight is live;
   - one-block-ahead makes block *N* and block *N+1* weights **live at the same time** — same key, two tensors in flight;
   - fix: a chain-position **parity (0,1,0,1…)** is folded into the cache key, so consecutive blocks alternate between **exactly two** buffers (counter cleared by `reset_gtp_state()`);
-  - without it the prefetch would **overwrite the weight the running GEMM is still reading** — a silent-correctness bug, not a crash.
+  - without it the prefetch would **overwrite the weight the running GEMM is still reading** — a silent-correctness bug, not a crash;
+  - the hazard is not exclusive to grouped chains — *any* chain whose neighbours share a key has it. Grouped chains are same-key throughout, so they take the blanket counter; others take the narrower `_ensure_distinct_buffer_from_prev` check, which allocates only where the collision is real (see *Buffer / memory management*).
 - **Eager only** — the optimization disables itself under CUDA-graph capture:
   - `_classify_param_chain` evaluates `graphed = _FULL_ITERATION or ("moe" in cuda_graph_modules)` **before** the split, and returns the plain `GRAPHED` chain when it is true;
   - so with `--cuda-graph-impl full_iteration` **every** param is `GRAPHED` — expert weights included — and they keep the ordinary one-step-ahead prefetch;

megatron/core/tensor_parallel/generalized_tensor_parallelism.py

@@ -1008,6 +1008,33 @@ def _double_buffer_parity(self) -> int:
             self._buf_parity = p
         return p
 
+    def _gather_buffer_identity(self, dtype) -> tuple:
+        """The part of the cache key that decides which weights share a gather buffer."""
+        return (self._unsharded_shape_padded, dtype, self.expert_idx)
+
+    def _ensure_distinct_buffer_from_prev(self, dtype):
+        """Move self to a second buffer if its chain predecessor would share one.
+
+        One-step-ahead prefetch keeps prev_w and self live at once, so sharing a buffer lets
+        self's gather clobber the weight prev_w's GEMM is still reading. Neighbours normally
+        differ in shape; a CUDA-graph-partitioned chain can leave two same-shaped weights
+        adjacent (embedding + output_layer alone in the UNGRAPHED chain).
+
+        Grouped chains use their own counter (``_GTP_GROUPED_BUF_PARITY_COUNTER``).
+        """
+        prev = self.prev_w
+        if prev is None or _chain_is_grouped(self.chain_id):
+            return
+        if self.is_routed_expert or prev.is_routed_expert:
+            return
+        if prev._cached_dtypes is None:  # never gathered — no buffer to collide with
+            return
+        if prev._gather_buffer_identity(prev._cached_dtypes[0]) != self._gather_buffer_identity(
+            dtype
+        ):
+            return
+        self._buf_parity = 1 - (getattr(prev, "_buf_parity", None) or 0)
+
     def _get_cache_key(self, dtype, fwd: bool, reduce_scatter: bool) -> tuple:
         """Build cache key from output shape + dtype.
 
@@ -1034,6 +1061,10 @@ def _get_cache_key(self, dtype, fwd: bool, reduce_scatter: bool) -> tuple:
             # chain_id keeps fc1/fc2 apart (both can be in flight at once, even if same-shaped);
             # parity alternates consecutive blocks between two buffers.
             key = key + (self.chain_id, self._double_buffer_parity())
+        elif getattr(self, "_buf_parity", None):
+            # Set by _ensure_distinct_buffer_from_prev. Parity 0 keeps the shared buffer, so
+            # only the second weight of an adjacent same-key pair costs an extra allocation.
+            key = key + (self._buf_parity,)
         return key
 
     def _strip_padding(self, tensor):
@@ -1432,6 +1463,9 @@ def all_gather_and_prefetch(self, fwd: bool = True, nvtx_label: str = None):
             dtypes = [
                 q.dtype if q is not None else w.dtype for q, w in zip(quantizers, self._weights)
             ]
+            # Must run before the reserve below — it decides which buffer the ticket gets.
+            self._ensure_distinct_buffer_from_prev(dtypes[0])
+
             for w, dt in zip(self._weights, dtypes):
                 w._ag_ticket_fwd = cache.reserve(w, dt, fwd=True)
                 cache.get(w._ag_ticket_fwd)

tests/unit_tests/generalized_tensor_parallel/test_gtp_basics.py

@@ -518,6 +518,75 @@ def test_parity_cached_and_stable(self):
         assert fwd[-1] == 0 and bwd[-1] == 0 and rs[-1] == 0
 
 
+def _worker_same_key_neighbours_dont_share_a_buffer(rank, world_size, port):
+    """Two same-shaped weights adjacent in a plain (non-grouped) chain need two buffers.
+
+    This is the GTP_ungraphed chain under CUDA graphs: every layer weight is captured, leaving
+    only embedding + output_layer behind — same shape, same dtype, same cache key. One-step-ahead
+    prefetch keeps both live (w0's consume issues w1's gather), so one buffer means w1's gather
+    lands on the weight w0's GEMM is still reading.
+    """
+    torch.manual_seed(0)
+    in_f, out_f = 32, 64
+    dtype = torch.bfloat16
+    gtp_remat_group = dist.new_group(list(range(world_size)))
+
+    l0 = _make_gtp_linear(in_f, out_f, gtp_remat_group, dtype)
+    l1 = _make_gtp_linear(in_f, out_f, gtp_remat_group, dtype)
+    # Distinct values so a clobber shows up in the gathered data, not just in the address.
+    with torch.no_grad():
+        l0.weight.fill_(1.0)
+        l1.weight.fill_(-1.0)
+
+    inp = torch.randn(4, in_f, dtype=dtype, device="cuda")
+    dist.broadcast(inp, src=0)
+
+    # First pass wires the chain; the second is the one that prefetches.
+    for _ in range(2):
+        l0(inp, is_first_microbatch=True)
+        l1(inp, is_first_microbatch=True)
+
+    w0, w1 = l0.weight, l1.weight
+    assert w0.next_w is w1 and w1.prev_w is w0, "w0 and w1 must be adjacent in one chain"
+    assert w0._gather_buffer_identity(dtype) == w1._gather_buffer_identity(
+        dtype
+    ), "nothing is being tested unless both weights want the same buffer"
+
+    cache = gtp_module.get_global_GTP_cache()
+    b0, b1 = cache.get(w0._ag_ticket_fwd), cache.get(w1._ag_ticket_fwd)
+    torch.cuda.synchronize()
+    assert b0.data_ptr() != b1.data_ptr(), (
+        "same-key chain neighbours share one gather buffer — w1's prefetch can clobber "
+        "the weight w0's GEMM is still reading"
+    )
+    assert (b0 == 1).all(), "w0's gathered weight was clobbered by w1's prefetch"
+    assert (b1 == -1).all(), "w1's buffer does not hold w1's weight"
+
+
+class TestSameKeyNeighbourTiebreak:
+    """A non-grouped chain buys a second buffer only where two adjacent weights would actually
+    share one — unlike the grouped chains above, which alternate unconditionally."""
+
+    _Fake = TestGroupedDoubleBuffer._Fake
+
+    def _key(self, parity=None):
+        f = self._Fake("GTP_ungraphed")
+        if parity is not None:
+            f._buf_parity = parity
+        return f._get_cache_key(torch.bfloat16, fwd=True, reduce_scatter=False)
+
+    def test_parity_zero_keeps_the_shared_buffer(self):
+        # Unset and 0 must give the same key: only the second weight of a pair pays.
+        assert self._key(0) == self._key() == ((128, 256), torch.bfloat16, 0, False)
+
+    def test_parity_one_gets_its_own_buffer(self):
+        assert self._key(1) != self._key()
+
+    def test_adjacent_same_shape_weights_get_distinct_buffers(self):
+        _requires_multi_gpu(4)
+        _run_distributed(_worker_same_key_neighbours_dont_share_a_buffer, 4)
+
+
 # ---------------------------------------------------------------------------
 # Wgrad reduce-scatter: shape and deferred async path
 # ---------------------------------------------------------------------------