Address PR review (items 2+3): reuse CP groups across pool cases

world_size and the rank set don't change for the lifetime of one pool, so
recreating the world group and a2a+p2p sub-groups per case wastes ~50-100 ms
of NCCL setup each. Pre-create them once in the pool worker (new helper
_create_cp_comm_groups), stash on the run_attention_with_cp module via
module-level _pool_cp_comm_group / _pool_cp_comm_sub_groups pointers, and
reuse them from run_dpa_with_cp in pool mode. Pool teardown destroys them
once at shutdown.

Also move per-case dist.new_group() calls inside the try/finally in
run_dpa_with_cp: a failure mid-loop in the a2a+p2p sub_group population
otherwise leaks every communicator created before the failure. The finally
now only destroys groups we created locally (cp_comm_group / sub_groups
populated in the else-branch), leaving pool-owned groups alone for reuse.

cyanguwa's review feedback on PR #2993.

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

Author: sudhakarsingh27

tests/pytorch/attention/run_attention_with_cp.py

@@ -30,6 +30,15 @@
 )
 from utils import ModelConfig, compare_and_assert
 
+# Pool mode (NVTE_CP_POOL_PG=1) only: shared CP collective groups, created once
+# per pool by run_attention_with_cp_pool.main() and reused across every case in
+# that pool. world_size and the rank set don't change per case, so re-creating
+# these per call would be wasted NCCL setup (~50-100 ms each). Single-shot
+# subprocess mode leaves these None / [] and run_dpa_with_cp creates/destroys
+# its own groups inline.
+_pool_cp_comm_group = None
+_pool_cp_comm_sub_groups: list = []
+
 dtypes = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp8": torch.bfloat16}
 
 
@@ -249,28 +258,36 @@ def run_dpa_with_cp(
     if not _pool_managed_pg:
         dist.init_process_group(backend="nccl", world_size=world_size, rank=rank)
 
-    # set up communication group for CP
+    # Set up communication group for CP. In pool mode, the pool worker has
+    # already pre-created world-scoped and a2a+p2p sub-groups once and stashed
+    # them in module-level pointers; we reuse those and the pool destroys them
+    # at shutdown. In single-shot mode we create them per call and destroy in
+    # the finally below.
     cp_comm_ranks = range(world_size)
     assert rank in cp_comm_ranks
-    cp_comm_group = dist.new_group(cp_comm_ranks, backend="nccl")
-    # Always defined so the finally cleanup below is safe even when cp_comm_type != "a2a+p2p".
-    cp_comm_sub_groups = []
-    if cp_comm_type == "a2a+p2p":
-        assert world_size % 2 == 0, (
-            "{cp_comm_type=} requires world_size % 2 = 0 as it assumes the a2a level has cp_size"
-            " = 2."
-        )
-        cp_comm_sub_ranks = [range(i * 2, (i + 1) * 2) for i in range(world_size // 2)]
-        cp_comm_sub_ranks += [range(i, world_size, 2) for i in range(2)]
-        for sub_ranks in cp_comm_sub_ranks:
-            sub_group = dist.new_group(sub_ranks, backend="nccl")
-            if rank in sub_ranks:
-                cp_comm_sub_groups.append(sub_group)
-
-    # Wrap test work in try/finally so the per-case NCCL communicators created above
-    # are destroyed even when the body raises (otherwise pool mode leaks one or more
-    # communicators per failed case, eventually exhausting NCCL resources).
+    _reusing_pool_groups = _pool_managed_pg and _pool_cp_comm_group is not None
+    cp_comm_group = None
+    cp_comm_sub_groups: list = []
+    # Wrap setup + body so any failure between new_group calls and the end of
+    # the test still hits the cleanup below — otherwise a partial sub-group
+    # population (e.g. NCCL refuses new_group mid-loop) leaks communicators.
     try:
+        if _reusing_pool_groups:
+            cp_comm_group = _pool_cp_comm_group
+            cp_comm_sub_groups = _pool_cp_comm_sub_groups if cp_comm_type == "a2a+p2p" else []
+        else:
+            cp_comm_group = dist.new_group(cp_comm_ranks, backend="nccl")
+            if cp_comm_type == "a2a+p2p":
+                assert world_size % 2 == 0, (
+                    "{cp_comm_type=} requires world_size % 2 = 0 as it assumes the a2a level has"
+                    " cp_size = 2."
+                )
+                cp_comm_sub_ranks = [range(i * 2, (i + 1) * 2) for i in range(world_size // 2)]
+                cp_comm_sub_ranks += [range(i, world_size, 2) for i in range(2)]
+                for sub_ranks in cp_comm_sub_ranks:
+                    sub_group = dist.new_group(sub_ranks, backend="nccl")
+                    if rank in sub_ranks:
+                        cp_comm_sub_groups.append(sub_group)
         if dtype == "fp8":
             if scaling_mode == "delayed":
                 fp8_recipe = DelayedScaling(fp8_dpa=fp8_dpa, fp8_mha=fp8_mha)
@@ -805,19 +822,24 @@ def run_dpa_with_cp(
                 logging.info(f"[Rank {rank}] CP vs no-CP: {names[i]} matches")
 
     finally:
-        # destroy distribution group(s); runs even on exception
-        if _pool_managed_pg:
-            # Pool owns the main PG; only clean up groups created for this case.
-            try:
-                dist.destroy_process_group(cp_comm_group)
-            except Exception:
-                pass
+        # Destroy only groups WE created. In pool mode with shared groups,
+        # cp_comm_group / cp_comm_sub_groups are owned by the pool runner and
+        # destroyed at pool shutdown — touching them here would tear down the
+        # cache that subsequent cases want to reuse.
+        if not _reusing_pool_groups:
+            if cp_comm_group is not None:
+                try:
+                    dist.destroy_process_group(cp_comm_group)
+                except Exception:
+                    pass
             for g in cp_comm_sub_groups:
                 try:
                     dist.destroy_process_group(g)
                 except Exception:
                     pass
-        else:
+        # Tear down the main PG only in single-shot mode. In pool mode the
+        # pool runner owns the main PG and destroys it at shutdown.
+        if not _pool_managed_pg:
             try:
                 dist.destroy_process_group()
             except Exception:

tests/pytorch/attention/run_attention_with_cp_pool.py

@@ -15,7 +15,7 @@
         call run_dpa_with_cp(**kwargs) — the same work function the
         per-case subprocess design uses, with NVTE_CP_POOL_PG=1 so the
         function reuses our PG instead of re-initing it
-        torch.cuda.empty_cache() + gc.collect() per case
+        torch.cuda.empty_cache() per case
     all ranks gather (ok, error_msg) to rank 0
     rank 0:
         write one JSON response line to stdout
@@ -26,7 +26,6 @@
     response: {"ok": true}
               {"ok": false, "error": "first failing rank's traceback"}
 """
-import gc
 import json
 import os
 import sys
@@ -43,18 +42,6 @@
 from transformer_engine.pytorch.quantization import FP8GlobalStateManager
 
 
-# Env vars run_dpa_with_cp re-sets at the top of every call. We pop them
-# defensively between cases so a future caller that *doesn't* re-set them
-# can't inherit a leftover value from a previous case in the same worker.
-_TRANSIENT_ENV_KEYS = (
-    "NVTE_FP8_DPA_BWD",
-    "NVTE_DPA_FP8CS_O_in_F16",
-    "NVTE_FLASH_ATTN",
-    "NVTE_FUSED_ATTN",
-    "NVTE_ALLOW_NONDETERMINISTIC_ALGO",
-)
-
-
 def _recv_request(rank: int) -> dict:
     box = [None]
     if rank == 0:
@@ -79,17 +66,17 @@ def _send_response(rank: int, payload: dict) -> None:
 def _reset_between_cases() -> None:
     """Drop state that would otherwise cascade across cases.
 
-    Matches the per-case startup of the single-shot worker (``_run_single_config``
-    on the per-case-subprocess branch): identical RNG seed at the start of every
-    case, FP8 state cleared, transient env vars cleared, allocator clean.
+    Matches the per-case startup of the single-shot worker
+    (``_run_single_config`` on the per-case-subprocess branch): identical RNG
+    seed at the start of every case, FP8 state cleared, allocator clean.
+    ``run_dpa_with_cp`` re-sets ``NVTE_FUSED_ATTN``/``NVTE_FLASH_ATTN``
+    unconditionally and pops the other transient env vars itself, so no
+    explicit pop is needed here.
     """
     torch.manual_seed(1234)
     torch.cuda.manual_seed(1234)
     FP8GlobalStateManager.reset()
-    for env_key in _TRANSIENT_ENV_KEYS:
-        os.environ.pop(env_key, None)
     torch.cuda.empty_cache()
-    gc.collect()
 
 
 _case_counter = 0
@@ -125,13 +112,49 @@ def _run_one(req: dict, rank: int) -> tuple[bool, str]:
     return ok, err
 
 
+def _create_cp_comm_groups(rank: int, world_size: int) -> tuple:
+    """Pre-create the CP collective groups for this pool.
+
+    world_size and the rank set are constant for the lifetime of one pool, so
+    the world group and the a2a+p2p sub-groups are deterministic. Creating
+    them once here and reusing them across every case eliminates ~50-100 ms
+    of NCCL setup per case (cyanguwa's review feedback on PR #2993).
+
+    Returns ``(world_group, a2a_p2p_sub_groups)``. ``a2a_p2p_sub_groups`` is
+    empty when world_size is too small to support a2a+p2p (needs an even
+    world_size ≥ 4); cases with cp_comm_type='a2a+p2p' wouldn't be routed to
+    such a pool anyway.
+    """
+    world_group = dist.new_group(range(world_size), backend="nccl")
+    sub_groups: list = []
+    if world_size >= 4 and world_size % 2 == 0:
+        # Mirror the layout in run_attention_with_cp.py: cp_size/2 pairs along
+        # axis 0, plus 2 stride-2 groups along axis 1.
+        cp_comm_sub_ranks = [range(i * 2, (i + 1) * 2) for i in range(world_size // 2)]
+        cp_comm_sub_ranks += [range(i, world_size, 2) for i in range(2)]
+        for sub_ranks in cp_comm_sub_ranks:
+            sub_group = dist.new_group(sub_ranks, backend="nccl")
+            if rank in sub_ranks:
+                sub_groups.append(sub_group)
+    return world_group, sub_groups
+
+
 def main() -> None:
     rank = int(os.environ["RANK"])
     world_size = int(os.environ["WORLD_SIZE"])
     torch.cuda.set_device(rank % torch.cuda.device_count())
     dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)
     os.environ["NVTE_CP_POOL_PG"] = "1"
 
+    # Stash pool-shared CP groups on the run_attention_with_cp module so
+    # run_dpa_with_cp can read them per case. Imported here (after the env var
+    # is set) to keep import-time side effects minimal.
+    import run_attention_with_cp as _rac
+
+    _rac._pool_cp_comm_group, _rac._pool_cp_comm_sub_groups = _create_cp_comm_groups(
+        rank, world_size
+    )
+
     try:
         while True:
             req = _recv_request(rank)
@@ -141,12 +164,10 @@ def main() -> None:
             ok, msg = _run_one(req, rank)
 
             gathered: list[tuple[bool, str]] = [None] * world_size  # type: ignore[list-item]
+            # gather_object is itself a collective synchronization point — if
+            # every rank reached it, none is ahead. No extra barrier needed.
             dist.gather_object((ok, msg), gathered if rank == 0 else None, dst=0)
 
-            # Surface a wedged communicator here, before the next case's
-            # collectives can inherit the corruption.
-            dist.barrier()
-
             if rank == 0:
                 all_ok = all(o for o, _ in gathered)
                 if all_ok:
@@ -155,6 +176,21 @@ def main() -> None:
                     first_err = next(m for o, m in gathered if not o)
                     _send_response(rank, {"ok": False, "error": first_err})
     finally:
+        # Tear down pool-shared CP groups before the main PG (NCCL requires
+        # sub-groups to be destroyed first). Each destroy is independently
+        # guarded so a wedged communicator on one group doesn't leak the rest.
+        if _rac._pool_cp_comm_group is not None:
+            try:
+                dist.destroy_process_group(_rac._pool_cp_comm_group)
+            except Exception:
+                pass
+        for g in _rac._pool_cp_comm_sub_groups:
+            try:
+                dist.destroy_process_group(g)
+            except Exception:
+                pass
+        _rac._pool_cp_comm_group = None
+        _rac._pool_cp_comm_sub_groups = []
         dist.destroy_process_group()