[PyTorch] Fix NCCL communicator init in cuSOLVERMp context creation

tests/pytorch/distributed/run_newton_schulz.py

@@ -77,15 +77,17 @@ def main():
     else:
         A = torch.empty(m, n, device="cuda", dtype=dtype)
 
-    # Broadcast the full matrix to all ranks
-    dist.broadcast(A, src=0)
-
-    # Scatter columns to each rank
-    local_cols = n // world_size
-    x_local = A[:, rank * local_cols : (rank + 1) * local_cols].contiguous()
-
+    # Context creation is the first collective, exercising lazy NCCL
+    # communicator materialization in CusolverMpCtx.
     ctx = CusolverMpCtx(dist.group.WORLD)
     try:
+        # Broadcast the full matrix to all ranks
+        dist.broadcast(A, src=0)
+
+        # Scatter columns to each rank
+        local_cols = n // world_size
+        x_local = A[:, rank * local_cols : (rank + 1) * local_cols].contiguous()
+
         newton_schulz(x_local, ctx, args.num_iterations, coefficients=coefficients)
     finally:
         ctx.destroy()

transformer_engine/common/newton_schulz/newton_schulz.cpp

@@ -134,6 +134,19 @@ void FreeWorkspace(NVTECusolverMpCtx* ctx) {
 
 NVTECusolverMpCtx* nvte_cusolvermp_ctx_create(ncclComm_t comm, int nranks, int rank) {
   NVTE_API_CALL(nvte_cusolvermp_ctx_create);
+  NVTE_CHECK(comm != nullptr, "NCCL communicator must be non-null");
+  NVTE_CHECK(nranks > 0, "Number of ranks must be positive, got ", nranks);
+  NVTE_CHECK(rank >= 0 && rank < nranks, "Rank ", rank, " is outside [0, ", nranks, ")");
+
+  int comm_nranks{};
+  int comm_rank{};
+  NVTE_CHECK_NCCL(ncclCommCount(comm, &comm_nranks));
+  NVTE_CHECK_NCCL(ncclCommUserRank(comm, &comm_rank));
+  NVTE_CHECK(comm_nranks == nranks, "NCCL communicator has ", comm_nranks,
+             " ranks, but the process group reports ", nranks);
+  NVTE_CHECK(comm_rank == rank, "NCCL communicator rank is ", comm_rank,
+             ", but the process group reports ", rank);
+
   int device_id{};
   NVTE_CHECK_CUDA(cudaGetDevice(&device_id));
 

transformer_engine/pytorch/newton_schulz.py

@@ -127,33 +127,59 @@ class CusolverMpCtx:
     """cuSolverMp context for Newton-Schulz matrix orthogonalization.
 
     Context creation is expensive; create once and reuse across multiple
-    :func:`newton_schulz` calls.  Call :meth:`destroy` when done.
+    :func:`newton_schulz` calls. Creation is collective over ``group`` and
+    must be called by every group member on its intended CUDA device. Call
+    :meth:`destroy` before destroying ``group``.
     """
 
     def __init__(self, group: dist.ProcessGroup) -> None:
+        # The cuSolverMp grid borrows the ProcessGroupNCCL communicator. Keep
+        # the group alive until the native context has released the grid.
+        self._ptr: Optional[int] = None
+        self._group: Optional[dist.ProcessGroup] = group
+
+        if not dist.is_initialized():
+            raise RuntimeError(
+                "torch.distributed must be initialized before creating CusolverMpCtx"
+            )
+
         self.nranks = dist.get_world_size(group)
-        self._ptr = tex.cusolvermp_ctx_create(
-            _get_nccl_comm_ptr(group), dist.get_world_size(group), dist.get_rank(group)
-        )
+        self.rank = dist.get_rank(group)
+        if self.rank < 0:
+            raise RuntimeError("The current process is not a member of the supplied process group")
+
+        comm_ptr = _get_nccl_comm_ptr(group)
+        self._ptr = tex.cusolvermp_ctx_create(comm_ptr, self.nranks, self.rank)
 
     def destroy(self) -> None:
         """Destroy the underlying cuSolverMp context."""
-        if self._ptr is not None:
-            tex.cusolvermp_ctx_destroy(self._ptr)
+        try:
+            if self._ptr is not None:
+                tex.cusolvermp_ctx_destroy(self._ptr)
+        finally:
             self._ptr = None
+            self._group = None
 
     def __del__(self) -> None:
         # Called when the context is manually destroyed or during Python teardown
         self.destroy()
 
 
 def _get_nccl_comm_ptr(group: dist.ProcessGroup) -> int:
-    """Extract the raw NCCL communicator pointer from a PyTorch process group."""
+    """Materialize and borrow a raw NCCL communicator from a process group."""
     backend = dist.get_backend(group)
     if backend != "nccl":
         raise RuntimeError(f"Newton-Schulz requires NCCL backend, got '{backend}'")
+
+    # ProcessGroupNCCL creates communicators lazily. This device-specific
+    # collective ensures that the communicator returned by _comm_ptr() exists
+    # and is ready on every group rank before cuSolverMp borrows it.
+    dist.barrier(group=group, device_ids=[torch.cuda.current_device()])
     nccl_backend = group._get_backend(torch.device("cuda"))
-    return nccl_backend._comm_ptr()
+    comm_ptr = nccl_backend._comm_ptr()
+    if not isinstance(comm_ptr, int) or comm_ptr == 0:
+        raise RuntimeError("ProcessGroupNCCL returned an invalid communicator pointer")
+    return comm_ptr
 
 
 def newton_schulz(