[Draft] Newton-Schulz via cuSOLVERMp

[vcherepanov-nv]

Description

Adds an API to call Newton-Schulz method on a distributed tensor.

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:

• Integrate cuSOLVERMp as a new dependency
• Add corresponding API to TE/common
• Add PyTorch binding and tests

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 adds distributed Newton-Schulz matrix orthogonalization to TransformerEngine via cuSolverMp integration. The implementation introduces a new optional dependency with proper build system integration, C API layer with RAII wrappers, and PyTorch bindings.

Key Changes

• Integrates cuSolverMp library as an optional dependency (enabled via NVTE_WITH_CUSOLVERMP)
• Adds C API (nvte_newton_schulz) with context management for cuSolverMp operations
• Provides Python API (te.pytorch.newton_schulz) for distributed tensor operations
• Includes distributed tests using torchrun with multiple GPUs

Issues Found

• Critical: Missing tensor contiguity validation - non-contiguous tensors will produce incorrect results since C++ code assumes contiguous memory layout
• Uses private PyTorch APIs (_get_backend, _comm_ptr) that may break in future PyTorch versions
• Test validates orthogonality (X @ X.t() ≈ I) but documentation claims "inverse square root" - semantic mismatch between implementation and documentation
• Context created/destroyed per call wastes resources - heavyweight operations (stream/event creation, cuSolverMp setup) happen on every invocation
• Row distribution assumes even splitting without validation - incorrect global dimension calculation if matrix rows not evenly divisible by number of ranks
• Synchronous cudaMalloc/cudaFree on hot path causes device synchronization, negating async benefits

Confidence Score: 2/5

• Not safe to merge - has a critical data correctness issue with missing contiguity validation
• The missing contiguity check is a critical bug that will cause incorrect results for non-contiguous tensors. Combined with multiple existing issues around performance, API stability, and semantic correctness, this PR needs significant fixes before merge.
• transformer_engine/pytorch/newton_schulz.py requires immediate attention for contiguity validation. Also review the test verification logic and documentation to clarify whether this computes inverse square root or orthogonalization.

Important Files Changed

Filename	Overview
transformer_engine/common/newton_schulz/newton_schulz.cpp	C++ implementation with proper RAII wrappers and error handling, but has performance concerns with context recreation and synchronous memory allocation
transformer_engine/pytorch/newton_schulz.py	Python API with missing contiguity validation (critical), uses private PyTorch APIs, and doesn't validate even row distribution
transformer_engine/pytorch/init.py	unconditionally imports optional feature - exposes newton_schulz even when built without NVTE_WITH_CUSOLVERMP
tests/pytorch/distributed/run_newton_schulz.py	distributed test worker with incorrect verification (checks orthogonality instead of inverse square root property)

Sequence Diagram

sequenceDiagram
    participant User as User Code
    participant PyAPI as newton_schulz.py
    participant Ext as C++ Extension
    participant Ctx as CusolverMpCtx
    participant cuSolver as cuSolverMp Library
    participant NCCL as NCCL Comm

    User->>PyAPI: newton_schulz(x, group, iterations)
    PyAPI->>PyAPI: Extract NCCL comm ptr from group
    PyAPI->>PyAPI: Compute global matrix dims (m = x.size(0) * nranks)
    PyAPI->>Ext: cusolvermp_ctx_create(comm_ptr, nranks, rank)
    Ext->>Ctx: Create context (stream, events, handle, grid)
    Ctx-->>Ext: ctx_ptr
    Ext-->>PyAPI: ctx_ptr
    
    PyAPI->>Ext: newton_schulz(ctx_ptr, m, n, x, iterations, coefficients)
    Ext->>Ext: Build NVTETensor from PyTorch tensor
    Ext->>Ctx: nvte_newton_schulz(ctx, m, n, x_tensor, ...)
    
    Note over Ctx: Event-based stream sync
    Ctx->>Ctx: Record event on caller_stream
    Ctx->>Ctx: Wait for event on internal stream
    
    Ctx->>Ctx: Compute ScaLAPACK distribution (mb, local_rows)
    Ctx->>cuSolver: cusolverMpNewtonSchulz_bufferSize(...)
    cuSolver-->>Ctx: workspace_size
    Ctx->>Ctx: Allocate/grow workspace (cudaMalloc)
    
    Ctx->>cuSolver: cusolverMpNewtonSchulz(...)
    cuSolver->>NCCL: Distributed matrix operations
    NCCL->>NCCL: Inter-rank communication
    cuSolver-->>Ctx: Result (in-place in x)
    
    Note over Ctx: Event-based stream sync
    Ctx->>Ctx: Record event on internal stream
    Ctx->>Ctx: Wait for event on caller_stream
    
    Ctx-->>Ext: Success
    Ext-->>PyAPI: Success
    
    PyAPI->>Ext: cusolvermp_ctx_destroy(ctx_ptr)
    Ext->>Ctx: Destroy context (free workspace, destroy handles)
    Ctx-->>Ext: Done
    Ext-->>PyAPI: Done
    PyAPI-->>User: Result (x modified in-place)

Loading

_{Last reviewed commit: 8eb6028}

[greptile-apps]

_{14 files reviewed, 12 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

verification doesn't match the comment - if X = A^{-1/2}, the check should be X @ A @ X ≈ I, not X @ X.t() ≈ I. The current check verifies X is orthogonal, not that X is the inverse square root of A. Note that A_orig is created on line 76 but never used.

Suggested change

	# Check: if X = A^{-1/2}, then X @ A @ X should be the identity matrix
	if rank == 0:
	XXT = X @ X.t()
	I = torch.eye(N, device=XXT.device, dtype=XXT.dtype)
	max_diff = (XXT - I).abs().max().item()
	print(f"Max |X @ X.t() - I|: {max_diff:.6e}", flush=True)
	# Check: if X = A^{-1/2}, then X @ A @ X should be the identity matrix
	XAX = X @ A_orig @ X
	I = torch.eye(N, device=XAX.device, dtype=XAX.dtype)
	max_diff = (XAX - I).abs().max().item()
	print(f"Max |X @ A @ X - I|: {max_diff:.6e}", flush=True)
	if torch.allclose(XAX, I, atol=args.atol, rtol=args.rtol):

[greptile-apps]

uses private PyTorch APIs (_get_backend, _comm_ptr) that may change in future versions

[greptile-apps]

coefficients mismatch with API defaults - test uses 15 coefficients for 5 iterations, but newton_schulz.py defaults to 5 coefficients. This inconsistency means default API behavior isn't tested.

[greptile-apps]

Documentation claims this computes "inverse square root" but the test validates orthogonality (X @ X.t() ≈ I), and commit dd1dd0b states "it approximates orthogonal matrix, not inverse square root". If this computes the polar decomposition (orthogonal factor), the documentation should be updated to reflect that. Inverse square root would satisfy X @ A @ X ≈ I, which is different from orthogonality.

[greptile-apps]

Assumes rows are evenly distributed (m = x.size(0) * nranks) but doesn't validate this. If matrix size isn't divisible by nranks, the computed global size m will be incorrect, leading to wrong results from cuSOLVERMp. Consider adding validation:

Suggested change

	m = x.size(0) * nranks # rows are distributed across ranks
	n = x.size(1)
	# Global matrix dimensions
	# Rows must be evenly distributed across ranks
	local_rows = x.size(0)
	m = local_rows * nranks
	n = x.size(1)

Then add a validation check that all ranks have the same local_rows via dist.all_reduce.

[greptile-apps]

Docstring says "inverse square root" but test checks orthogonality. Update to match actual behavior (see comment on header file).

[greptile-apps]

num_coefficients parameter is unused

The num_coefficients parameter is accepted but never referenced in the function body. Neither cusolverMpNewtonSchulz_bufferSize (line 166) nor cusolverMpNewtonSchulz (line 183) receive this value. If cuSolverMp infers the count from num_iterations internally, then num_coefficients is dead code that should be removed from the API. If cuSolverMp actually needs it, then it should be passed to the cuSolverMp calls — otherwise the library may read out of bounds on the coefficients array.

[greptile-apps]

Unconditional import of optional feature

newton_schulz is unconditionally imported and exported as part of the public API, even when TE is built without NVTE_WITH_CUSOLVERMP. While the function itself raises a runtime error when called, this exposes the symbol to all users and makes it appear as a supported feature in auto-complete and docs. Consider guarding this import behind a check (similar to how other optional features are handled), or at minimum adding a note in the docstring that the function requires NVTE_WITH_CUSOLVERMP=1 at build time.

_{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]

Synchronous cudaFree/cudaMalloc on hot path

cudaFree followed by cudaMalloc inside nvte_newton_schulz will synchronize with the device each time the workspace needs to grow. Since the context is recreated on every call from newton_schulz.py (line 82-86 creates + destroys ctx each invocation), the workspace will never be reused across calls — the grow-only caching here is ineffective. Consider either:

1. Allowing callers to keep the context alive across calls, or
2. Using cudaMallocAsync/cudaFreeAsync on ctx->stream to avoid synchronous stalls.

[greptile-apps]

Context created/destroyed per call wastes resources

A new NVTECusolverMpCtx is created and destroyed on every invocation of newton_schulz. Context creation involves cudaStreamCreate, two cudaEventCreate calls, cusolverMpCreate, and cusolverMpCreateDeviceGrid — all of which are heavyweight operations. And since the context is destroyed afterward, the grow-only workspace caching in the C++ layer (lines 170-177 of newton_schulz.cpp) is never actually reused.

Consider caching the context (e.g., in a module-level dict keyed by (nccl_comm_ptr, nranks, rank)) and reusing it across calls, or exposing the context lifecycle to callers so they can amortize the cost when calling newton_schulz repeatedly in a training loop.

[greptile-apps]

use ValueError instead of assert for validation - assert can be disabled with Python's -O flag

Suggested change

	assert (
	len(coefficients) == num_iterations * 3
	), f"Unexpected number of coefficients: {len(coefficients)} for {num_iterations} iterations"
	if len(coefficients) != num_iterations * 3:
	raise ValueError(
	f"Unexpected number of coefficients: {len(coefficients)} for {num_iterations} iterations"
	)

[greptile-apps]

missing contiguity check - C++ code uses data_ptr() which requires contiguous memory. Non-contiguous tensors will cause incorrect results.

Suggested change

	if x.dim() != 2:
	raise ValueError(f"Expected 2D tensor, got {x.dim()}D")
	if not x.is_cuda:
	raise ValueError("Input tensor must be on CUDA device")
	if x.dim() != 2:
	raise ValueError(f"Expected 2D tensor, got {x.dim()}D")
	if not x.is_cuda:
	raise ValueError("Input tensor must be on CUDA device")
	if not x.is_contiguous():
	raise ValueError("Input tensor must be contiguous")