From 61ec52f16981ef07fe88219f1cc6c4c95713db78 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Wed, 22 Jul 2026 21:00:37 +0000 Subject: [PATCH 1/2] test: cover lazy NCCL setup for Newton-Schulz Construct the cuSOLVERMp context before any other distributed collective so the distributed test exercises lazy ProcessGroupNCCL communicator initialization. Signed-off-by: Vladimir Cherepanov --- tests/pytorch/distributed/run_newton_schulz.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/tests/pytorch/distributed/run_newton_schulz.py b/tests/pytorch/distributed/run_newton_schulz.py index 712d83bd1c..9dbb323f09 100644 --- a/tests/pytorch/distributed/run_newton_schulz.py +++ b/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() From 29ee264f9c8207b94a474d3e144fa45ce471c3e6 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Wed, 22 Jul 2026 21:00:49 +0000 Subject: [PATCH 2/2] fix: initialize borrowed NCCL comm for Newton-Schulz Materialize the ProcessGroupNCCL communicator before borrowing its raw handle, retain the process group for the context lifetime, and validate the communicator size and rank during native context creation. Signed-off-by: Vladimir Cherepanov --- .../common/newton_schulz/newton_schulz.cpp | 13 ++++++ transformer_engine/pytorch/newton_schulz.py | 42 +++++++++++++++---- 2 files changed, 47 insertions(+), 8 deletions(-) diff --git a/transformer_engine/common/newton_schulz/newton_schulz.cpp b/transformer_engine/common/newton_schulz/newton_schulz.cpp index 0d6426a156..5eeaf2da00 100644 --- a/transformer_engine/common/newton_schulz/newton_schulz.cpp +++ b/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)); diff --git a/transformer_engine/pytorch/newton_schulz.py b/transformer_engine/pytorch/newton_schulz.py index 1cbe6ebfbf..852739f769 100644 --- a/transformer_engine/pytorch/newton_schulz.py +++ b/transformer_engine/pytorch/newton_schulz.py @@ -127,20 +127,38 @@ 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 @@ -148,12 +166,20 @@ def __del__(self) -> None: 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(