[PyTorch] Fix NCCL communicator init in cuSOLVERMp context creation

[vcherepanov-nv]

Description

Fix cuSOLVERMp context initialization when borrowing an NCCL communicator from a PyTorch process group.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Materializes the lazily initialized NCCL communicator before accessing _comm_ptr().
• Retains the process group for the cuSOLVERMp context lifetime.
• Validates the communicator size and rank in native context creation.
• Updates the distributed test so context creation is the first collective, covering lazy communicator initialization.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

This PR fixes a race condition where CusolverMpCtx could call _comm_ptr() on a ProcessGroupNCCL communicator that had not yet been materialized, because PyTorch creates NCCL communicators lazily on first use. The fix inserts a dist.barrier inside _get_nccl_comm_ptr to force communicator creation before the pointer is borrowed, retains a strong reference to the ProcessGroup for the context's lifetime, and adds C++-side validation that the borrowed communicator's rank/count match the Python-side values.

• newton_schulz.py: _get_nccl_comm_ptr now runs a device-specific barrier to materialize the NCCL communicator; CusolverMpCtx holds self._group to prevent premature GC; destroy() uses try/finally for unconditional cleanup.
• newton_schulz.cpp: nvte_cusolvermp_ctx_create gains null/bounds checks on comm, nranks, and rank, then cross-validates them against ncclCommCount/ncclCommUserRank.
• run_newton_schulz.py: CusolverMpCtx is created before dist.broadcast so context initialization is the first collective, specifically exercising the lazy-materialization path.

Confidence Score: 4/5

The fix is logically correct and addresses a real crash scenario; the only open questions are around private PyTorch API stability and cleanup safety during interpreter teardown.

The barrier-based materialization approach is sound, the _group retention correctly prevents the communicator from being freed while the cuSolverMp grid holds a raw pointer to it, and the C++ validation catches misconfiguration early. The reliance on _get_backend()._comm_ptr() (private PyTorch internals) is the main fragility.

transformer_engine/pytorch/newton_schulz.py — specifically the _get_nccl_comm_ptr private-API usage and the del teardown path.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/newton_schulz.py	Core fix: adds dist.barrier() in _get_nccl_comm_ptr() to force lazy NCCL communicator materialization, retains self._group reference, and hardens destroy() with try/finally. Logic is correct but relies on private PyTorch APIs (_get_backend, _comm_ptr).
transformer_engine/common/newton_schulz/newton_schulz.cpp	Adds defensive input validation (null comm, rank/nranks bounds) and cross-validates the NCCL communicator's actual count/rank against the Python-supplied values using ncclCommCount/ncclCommUserRank.
tests/pytorch/distributed/run_newton_schulz.py	Moves CusolverMpCtx creation before dist.broadcast to make context init the first collective, covering the lazy-communicator path. Wraps the body in try/finally for deterministic cleanup.

Comments Outside Diff (1)

1. tests/pytorch/distributed/run_newton_schulz.py, line 95-98 (link)

   x_local referenced outside the try block that defines it

   x_local is assigned inside the try block (line 89). If an exception escapes the try — e.g. dist.broadcast fails — x_local is never bound, yet the all_gather call on line 97 would raise a NameError, masking the original exception. In the happy path this is harmless, but it makes failure diagnosis harder in CI. Moving x_local initialization before the try block (similar to the original layout where it preceded CusolverMpCtx) or catching/re-raising would keep error messages clear.

_{Reviews (1): Last reviewed commit: "fix: initialize borrowed NCCL comm for N..." | Re-trigger Greptile}

[greptile-apps]

Private PyTorch API dependency

_get_backend() and _comm_ptr() are private ProcessGroupNCCL methods and can be renamed, replaced, or removed across PyTorch minor releases without a deprecation cycle. If the API changes, the failure mode is a silent AttributeError at context-creation time rather than a compilation error, which may only surface in distributed runs. It would be worth tracking the PyTorch version that introduced these methods and adding an assertion or a version guard, or at minimum filing a PyTorch issue to request a stable public surface for borrowing the raw ncclComm_t.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[greptile-apps]

__del__ may race with CUDA context teardown

__del__ calls destroy() which in turn calls tex.cusolvermp_ctx_destroy. During interpreter shutdown, CUDA driver state is torn down in an unspecified order relative to Python object finalization. If the CUDA context has already been destroyed when __del__ fires, cusolvermp_ctx_destroy (and the cusolverMpDestroyGrid / cusolverMpDestroy calls inside it) may trigger CUDA errors or a segfault. Guarding with try/except Exception: pass inside __del__, or checking torch.cuda.is_initialized() before calling into the native layer, would prevent noisy teardown-time crashes in programs that forget to call destroy() explicitly.