[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 integrates cuSolverMp to provide distributed Newton-Schulz matrix orthogonalization. The implementation adds ~650 lines across build configuration, C++ bindings, Python API, and tests.

Key changes:

• Adds cuSolverMp as optional dependency with auto-detection from common library symbols
• Implements C API wrapper in transformer_engine/common/newton_schulz/
• Provides PyTorch binding with context management and stream synchronization
• Includes distributed tests for float32/bfloat16 with 5 and 15 iterations

Major concerns from prior review rounds:

• Documentation and tests claim "inverse square root" but implementation computes orthogonal matrix (polar decomposition)
• Uses private PyTorch APIs (_get_backend, _comm_ptr) subject to breakage
• Creates/destroys heavyweight context on every call instead of caching
• Test verification checks wrong property (orthogonality vs inverse square root)
• Missing validation for even row distribution across ranks
• num_coefficients parameter validated but never used by cuSolverMp
• Synchronous cudaMalloc/cudaFree on hot path cause device synchronization

Additional issue found:

• Missing dtype validation despite documented float32/bfloat16 requirement

Confidence Score: 2/5

• Multiple functional and performance issues need resolution before merge
• Score reflects accumulated issues from multiple review rounds: incorrect test verification, misleading documentation (inverse square root vs orthogonalization), performance problems (per-call context creation, synchronous CUDA allocations), API stability risks (private PyTorch APIs), and missing validation (dtype, tensor distribution). While the build integration is solid, the core functionality has correctness and design issues that must be addressed.
• Primary attention needed on transformer_engine/pytorch/newton_schulz.py (API design, validation), tests/pytorch/distributed/run_newton_schulz.py (incorrect verification), and transformer_engine/common/newton_schulz/newton_schulz.cpp (performance optimizations)

Important Files Changed

Filename	Overview
transformer_engine/common/newton_schulz/newton_schulz.cpp	Core C++ implementation using cuSolverMp. Previous reviews identified synchronous cudaMalloc on hot path and unused num_coefficients parameter (only validated, never passed to cuSolverMp)
transformer_engine/pytorch/newton_schulz.py	Python API with several issues: uses private PyTorch APIs, creates/destroys context per call (performance issue), missing dtype validation, and documents "inverse square root" but actually computes orthogonal matrix
tests/pytorch/distributed/run_newton_schulz.py	Test worker with incorrect verification - checks orthogonality (X @ X.t() ≈ I) instead of inverse square root (X @ A @ X ≈ I), and coefficient count mismatch with API defaults
transformer_engine/common/include/transformer_engine/newton_schulz.h	Header file with misleading documentation claiming "inverse square root" when it actually computes orthogonal matrix (polar decomposition)
transformer_engine/pytorch/init.py	Unconditionally imports newton_schulz even when feature not built, exposing it in public API despite potential runtime errors

Sequence Diagram

sequenceDiagram
    participant User as Python User
    participant PyAPI as newton_schulz.py
    participant PyExt as C++ Extension
    participant Common as newton_schulz.cpp
    participant cuSolver as cuSolverMp Library
    participant NCCL as NCCL Communicator

    User->>PyAPI: newton_schulz(x, group, iterations, coeffs)
    PyAPI->>PyAPI: Extract NCCL comm from ProcessGroup
    PyAPI->>PyAPI: Calculate global dims (m, n)
    PyAPI->>PyExt: cusolvermp_ctx_create(comm, nranks, rank)
    PyExt->>Common: nvte_cusolvermp_ctx_create()
    Common->>Common: Create CUDA stream & events
    Common->>cuSolver: cusolverMpCreate()
    Common->>cuSolver: cusolverMpCreateDeviceGrid()
    Common-->>PyExt: Return context pointer
    PyExt-->>PyAPI: Return context handle
    
    PyAPI->>PyExt: newton_schulz(ctx, m, n, x, iters, coeffs)
    PyExt->>Common: nvte_newton_schulz()
    Common->>Common: Stream synchronization (events)
    Common->>cuSolver: cusolverMpNewtonSchulz_bufferSize()
    cuSolver-->>Common: Workspace size
    Common->>Common: Allocate/grow workspace (cudaMalloc)
    Common->>cuSolver: cusolverMpNewtonSchulz()
    cuSolver->>NCCL: Distributed matrix operations
    NCCL-->>cuSolver: Sync results
    cuSolver-->>Common: Modified matrix (in-place)
    Common->>Common: Stream synchronization (events)
    Common-->>PyExt: Success
    PyExt-->>PyAPI: Success
    
    PyAPI->>PyExt: cusolvermp_ctx_destroy(ctx)
    PyExt->>Common: nvte_cusolvermp_ctx_destroy()
    Common->>Common: Free workspace
    Common->>cuSolver: Destroy grid & handle
    Common->>Common: Destroy stream & events
    Common-->>PyExt: Done
    PyExt-->>PyAPI: Done
    PyAPI-->>User: Modified tensor x

Loading

_{Last reviewed commit: d3740fb}

[greptile-apps]

_{15 files reviewed, 13 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")

[greptile-apps]

Missing dtype validation - docstring on line 36 states tensor must be float32 or bfloat16, but this isn't enforced. Passing unsupported dtypes leads to confusing errors from cuSolverMp.

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 x.dtype not in (torch.float32, torch.bfloat16):
	raise ValueError(f"Input tensor must be float32 or bfloat16, got {x.dtype}")