From 3abd8728e41cdfee367d10d49c94314f031882da Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Thu, 17 Sep 2026 20:42:54 +0000 Subject: [PATCH 1/3] Fix distributed Newton-Schulz correctness - Explicitly enable cuSOLVERMp input normalization and compute-type Gram reduction. - Validate distributed dimensions and reject unsupported tall low-level inputs. - Reduce device workspace requirements to the grid-wide maximum. - Synchronize workspace cleanup and retain host workspace for asynchronous operations. - Normalize the test reference and evaluate Gram and reference errors in FP32. - Use rank-independent aligned test shapes and strict absolute tolerances. Tests: - pytest -q tests/pytorch/distributed/test_newton_schulz.py - torchrun workers with 1, 2, 4, and 8 processes Signed-off-by: Vladimir Cherepanov --- .../pytorch/distributed/run_newton_schulz.py | 45 +++++++++--------- .../transformer_engine/newton_schulz.h | 2 +- .../common/newton_schulz/newton_schulz.cpp | 47 +++++++++++++++++-- .../pytorch/optimizers/newton_schulz.py | 7 ++- 4 files changed, 73 insertions(+), 28 deletions(-) diff --git a/tests/pytorch/distributed/run_newton_schulz.py b/tests/pytorch/distributed/run_newton_schulz.py index bd061949ad7..2724841b563 100644 --- a/tests/pytorch/distributed/run_newton_schulz.py +++ b/tests/pytorch/distributed/run_newton_schulz.py @@ -31,8 +31,9 @@ def newton_schulz_reference( in_x: torch.Tensor, coefficients: list[tuple[float, float, float]] ) -> torch.Tensor: - """Local Newton-Schulz reference mirroring the provided Octave update.""" - x = in_x.clone() + """Local Newton-Schulz reference matching cuSOLVERMp input normalization.""" + x = in_x.float().clone() + x /= torch.linalg.vector_norm(x) for a, b, c in coefficients: xxt = x @ x.mT x = a * x + b * xxt @ x + c * xxt @ xxt @ x @@ -47,30 +48,28 @@ def _dtype_from_name(dtype: str) -> torch.dtype: raise ValueError(f"Unsupported dtype: {dtype}") -def _test_tolerances(dtype: str, check: str, world_size: int) -> tuple[float, float]: - if dtype == "bfloat16": - return (5e-2, 5e-2) - if check == "orthogonality" and world_size == 1: - return (2e-2, 2e-2) - return (1e-2, 1e-2) +def _test_tolerances(dtype: str, check: str) -> tuple[float, float]: + if check == "reference": + return (1e-2 if dtype == "bfloat16" else 1e-3, 0.0) + return (1.5e-2 if dtype == "bfloat16" else 1e-2, 0.0) -def _shape_scale(world_size: int) -> int: - return 4 if world_size == 1 else world_size +def _aligned_size(size: int, world_size: int) -> int: + """Round a global distributed dimension up to a whole number of shards.""" + return (size + world_size - 1) // world_size * world_size def _orthogonality_shapes(world_size: int) -> list[tuple[int, int]]: - scale = _shape_scale(world_size) + rows = _aligned_size(512, world_size) return [ - (scale * 64, scale * 64), - (scale * 64, scale * 96), - (scale * 96, scale * 64), + (rows, rows), + (rows, _aligned_size(768, world_size)), ] def _reference_shapes(world_size: int) -> list[tuple[int, int]]: - scale = _shape_scale(world_size) - return [(scale * 64, scale * 64)] + size = _aligned_size(512, world_size) + return [(size, size)] def _make_matrix( @@ -117,7 +116,7 @@ def _run_case( dtype = _dtype_from_name(dtype_name) m, n = matrix_shape coefficients = get_coefficients(num_iterations, coeff_type) - atol, rtol = _test_tolerances(dtype_name, check, world_size) + atol, rtol = _test_tolerances(dtype_name, check) if api == "tp" and partition_dim is None: # Replicated inputs are sharded along the larger dimension for cuSolverMp. @@ -169,21 +168,23 @@ def _run_case( # Check: the resulting matrix should be orthogonal, or match a local reference. if check == "orthogonality": + X_float = X.float() if m <= n: - gram = X @ X.t() + gram = X_float @ X_float.t() expected = torch.eye(m, device=gram.device, dtype=gram.dtype) label = "X @ X.t() - I" else: - gram = X.t() @ X + gram = X_float.t() @ X_float expected = torch.eye(n, device=gram.device, dtype=gram.dtype) label = "X.t() @ X - I" max_diff = (gram - expected).abs().max().item() passed = torch.allclose(gram, expected, atol=atol, rtol=rtol) elif check == "reference": - reference = newton_schulz_reference(A.float(), coefficients).to(dtype) - max_diff = (X - reference).abs().max().item() + reference = newton_schulz_reference(A, coefficients) + X_float = X.float() + max_diff = (X_float - reference).abs().max().item() label = "distributed - reference" - passed = torch.allclose(X, reference, atol=atol, rtol=rtol) + passed = torch.allclose(X_float, reference, atol=atol, rtol=rtol) else: raise ValueError(f"Unsupported check: {check}") diff --git a/transformer_engine/common/include/transformer_engine/newton_schulz.h b/transformer_engine/common/include/transformer_engine/newton_schulz.h index bea8e32b1ef..f4d22ba935e 100644 --- a/transformer_engine/common/include/transformer_engine/newton_schulz.h +++ b/transformer_engine/common/include/transformer_engine/newton_schulz.h @@ -45,7 +45,7 @@ void nvte_cusolvermp_ctx_destroy(NVTECusolverMpCtx* ctx); /*! \brief Compute Newton-Schulz matrix orthogonalization in-place. * * \param[in] ctx cuSolverMp context. - * \param[in] m Global number of rows. + * \param[in] m Global number of rows. Must be no greater than n. * \param[in] n Global number of columns. * \param[in,out] x Local part of the matrix (modified in-place). * \param[in] num_iterations Number of Newton-Schulz iterations. diff --git a/transformer_engine/common/newton_schulz/newton_schulz.cpp b/transformer_engine/common/newton_schulz/newton_schulz.cpp index 5eeaf2da005..a6af4a1f3bc 100644 --- a/transformer_engine/common/newton_schulz/newton_schulz.cpp +++ b/transformer_engine/common/newton_schulz/newton_schulz.cpp @@ -8,6 +8,7 @@ #include +#include #include #include @@ -103,6 +104,7 @@ CudaEvent MakeCudaEvent() { struct NVTECusolverMpCtx { int64_t nranks; int64_t rank; + ncclComm_t comm; CudaStream stream; CudaEvent in_ready; CudaEvent out_ready; @@ -111,11 +113,14 @@ struct NVTECusolverMpCtx { void* workspace; size_t workspace_size; bool workspace_registered; + std::vector workspace_host; }; namespace { void FreeWorkspace(NVTECusolverMpCtx* ctx) { + // Buffer deregistration and grid destruction require all grid work to be complete. + NVTE_CHECK_CUDA(cudaStreamSynchronize(ctx->stream.get())); if (ctx->workspace == nullptr) { return; } @@ -130,6 +135,25 @@ void FreeWorkspace(NVTECusolverMpCtx* ctx) { ctx->workspace_registered = false; } +size_t GridMaxWorkspaceSize(NVTECusolverMpCtx* ctx, size_t local_size) { + if (ctx->nranks == 1) { + return local_size; + } + + uint64_t size = local_size; + uint64_t* device_size = nullptr; + NVTE_CHECK_CUDA(cudaMalloc(&device_size, sizeof(size))); + NVTE_CHECK_CUDA( + cudaMemcpyAsync(device_size, &size, sizeof(size), cudaMemcpyHostToDevice, ctx->stream.get())); + NVTE_CHECK_NCCL(ncclAllReduce(device_size, device_size, 1, ncclUint64, ncclMax, ctx->comm, + ctx->stream.get())); + NVTE_CHECK_CUDA( + cudaMemcpyAsync(&size, device_size, sizeof(size), cudaMemcpyDeviceToHost, ctx->stream.get())); + NVTE_CHECK_CUDA(cudaStreamSynchronize(ctx->stream.get())); + NVTE_CHECK_CUDA(cudaFree(device_size)); + return static_cast(size); +} + } // namespace NVTECusolverMpCtx* nvte_cusolvermp_ctx_create(ncclComm_t comm, int nranks, int rank) { @@ -160,6 +184,7 @@ NVTECusolverMpCtx* nvte_cusolvermp_ctx_create(ncclComm_t comm, int nranks, int r return new NVTECusolverMpCtx{ nranks, rank, + comm, std::move(stream), std::move(in_ready), std::move(out_ready), @@ -168,6 +193,7 @@ NVTECusolverMpCtx* nvte_cusolvermp_ctx_create(ncclComm_t comm, int nranks, int r nullptr, 0, false, + {}, }; } @@ -187,6 +213,7 @@ void nvte_newton_schulz(NVTECusolverMpCtx* ctx, int64_t m, int64_t n, NVTETensor NVTE_CHECK(num_coefficients == num_iterations * 3, num_iterations, " iterations require ", num_iterations * 3, " coefficients, but ", num_coefficients, " are passed"); const auto* t = convertNVTETensorCheck(x); + NVTE_CHECK(m <= n, "Column-sharded Newton-Schulz requires rows <= columns, got ", m, " > ", n); // Make the internal stream wait for the caller's stream so that // the input tensor is ready before cuSolverMp reads it. @@ -200,6 +227,7 @@ void nvte_newton_schulz(NVTECusolverMpCtx* ctx, int64_t m, int64_t n, NVTETensor // Compute local leading dimension const int64_t local_cols = cusolverMpNUMROC(n, nb, ctx->rank, 0, ctx->nranks); NVTE_CHECK(t->shape().size() == 2, "Shape size:", t->shape().size()); + NVTE_CHECK(t->shape()[0] == m, "Tensor rows:", t->shape()[0], "Expected rows:", m); NVTE_CHECK(t->shape()[1] == local_cols, "Tensor cols:", t->shape()[1], "Local cols:", local_cols); const int64_t lld = std::max(local_cols, static_cast(1)); @@ -210,6 +238,13 @@ void nvte_newton_schulz(NVTECusolverMpCtx* ctx, int64_t m, int64_t n, NVTETensor // Create Newton-Schulz descriptor auto ns_desc = MakeCusolverMpNSDesc(); + const int enabled = 1; + NVTE_CHECK_CUSOLVERMP(cusolverMpNewtonSchulzDescriptorSetAttribute( + ns_desc.get(), CUSOLVERMP_NEWTON_SCHULZ_DESCRIPTOR_ATTRIBUTE_NORMALIZE, &enabled, + sizeof(enabled))); + NVTE_CHECK_CUSOLVERMP(cusolverMpNewtonSchulzDescriptorSetAttribute( + ns_desc.get(), CUSOLVERMP_NEWTON_SCHULZ_DESCRIPTOR_ATTRIBUTE_REDUCE_VIA_COMPUTE_TYPE, + &enabled, sizeof(enabled))); // Query workspace sizes size_t wrksp_size_device = 0; @@ -217,6 +252,7 @@ void nvte_newton_schulz(NVTECusolverMpCtx* ctx, int64_t m, int64_t n, NVTETensor NVTE_CHECK_CUSOLVERMP(cusolverMpNewtonSchulz_bufferSize( ctx->handle.get(), ns_desc.get(), n, m, t->data.dptr, 1, 1, mat_desc.get(), num_iterations, coefficients, CUDA_R_32F, &wrksp_size_device, &wrksp_size_host)); + wrksp_size_device = GridMaxWorkspaceSize(ctx, wrksp_size_device); // Allocate/grow device workspace if (ctx->workspace_size < wrksp_size_device) { @@ -244,14 +280,17 @@ void nvte_newton_schulz(NVTECusolverMpCtx* ctx, int64_t m, int64_t n, NVTETensor ctx->workspace_registered = workspace_registered; } - // Allocate host workspace - std::vector workspace_host(wrksp_size_host); + // Keep host workspace alive until all work on the internal stream is complete. + if (ctx->workspace_host.size() < wrksp_size_host) { + NVTE_CHECK_CUDA(cudaStreamSynchronize(ctx->stream.get())); + ctx->workspace_host.resize(wrksp_size_host); + } // Execute Newton-Schulz NVTE_CHECK_CUSOLVERMP(cusolverMpNewtonSchulz( ctx->handle.get(), ns_desc.get(), n, m, t->data.dptr, 1, 1, mat_desc.get(), num_iterations, - coefficients, CUDA_R_32F, ctx->workspace, ctx->workspace_size, workspace_host.data(), - workspace_host.size(), nullptr)); + coefficients, CUDA_R_32F, ctx->workspace, ctx->workspace_size, ctx->workspace_host.data(), + ctx->workspace_host.size(), nullptr)); // Make the caller's stream wait for the internal stream so that // the output tensor is ready before the caller uses it. diff --git a/transformer_engine/pytorch/optimizers/newton_schulz.py b/transformer_engine/pytorch/optimizers/newton_schulz.py index 4f868f1c127..f837627536e 100644 --- a/transformer_engine/pytorch/optimizers/newton_schulz.py +++ b/transformer_engine/pytorch/optimizers/newton_schulz.py @@ -202,7 +202,8 @@ def newton_schulz( x : torch.Tensor Local part of the distributed matrix (modified in-place). Must be a 2D CUDA tensor of type float32 or bfloat16. - Columns are distributed across ranks. + Columns are distributed across ranks. The global matrix must have no + more rows than columns; use :func:`newton_schulz_tp` for tall matrices. ctx : CusolverMpCtx cuSolverMp context created by :func:`cusolvermp_ctx_create`. num_iterations : int, optional @@ -237,6 +238,10 @@ def newton_schulz( # Global matrix dimensions; columns are distributed across ranks. m = x.size(0) n = x.size(1) * ctx.nranks + if m > n: + raise ValueError( + f"Expected global rows <= columns for the column-sharded API, got {m} > {n}" + ) tex.newton_schulz(ctx._ptr, m, n, x, num_iterations, flat_coefficients) From 33b711b11e6bd181db300fcb2ef70a8c9fe1f30b Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Fri, 18 Sep 2026 19:45:31 +0000 Subject: [PATCH 2/3] Cache distributed Newton-Schulz workspaces Cache cuSOLVERMp workspace configurations and retain the NCCL reduction scalar so repeated optimizer steps avoid redundant workspace queries, allocations, collectives, and stream synchronization. Restore tall tensor-parallel coverage for supported layouts and reject distributed partitions along the smaller matrix dimension with a clear error. Tests: - ninja -C build/cmake transformer_engine - pytest -q tests/pytorch/distributed/test_newton_schulz.py (2 passed) Signed-off-by: Vladimir Cherepanov --- .../pytorch/distributed/run_newton_schulz.py | 47 ++++++ .../common/newton_schulz/newton_schulz.cpp | 142 +++++++++++++----- .../pytorch/optimizers/newton_schulz.py | 17 ++- 3 files changed, 164 insertions(+), 42 deletions(-) diff --git a/tests/pytorch/distributed/run_newton_schulz.py b/tests/pytorch/distributed/run_newton_schulz.py index 2724841b563..2a38bc95f99 100644 --- a/tests/pytorch/distributed/run_newton_schulz.py +++ b/tests/pytorch/distributed/run_newton_schulz.py @@ -72,6 +72,10 @@ def _reference_shapes(world_size: int) -> list[tuple[int, int]]: return [(size, size)] +def _tall_reference_shape(world_size: int) -> tuple[int, int]: + return (_aligned_size(768, world_size), _aligned_size(512, world_size)) + + def _make_matrix( m: int, n: int, @@ -255,6 +259,49 @@ def run_all_tests(ctx: CusolverMpCtx) -> None: tp_mode=tp_mode, ) + tall_shape = _tall_reference_shape(world_size) + tall_tp_configs = ( + (0, "distributed"), + (0, "duplicated"), + (1, "duplicated"), + (None, "duplicated"), + ) + for partition_dim, tp_mode in tall_tp_configs: + config = (tall_shape, partition_dim, tp_mode) + if rank == 0: + print(f"Running tall TP API reference check with {config=}", flush=True) + _run_case( + ctx=ctx, + check="reference", + dtype_name="float32", + matrix_shape=tall_shape, + num_iterations=5, + coeff_type="quintic", + api="tp", + partition_dim=partition_dim, + tp_mode=tp_mode, + ) + + # A directly column-sharded tall matrix cannot be transposed into the column distribution + # required by the low-level API without first redistributing it. + m, n = tall_shape + x_local = torch.empty(m, n // world_size, device="cuda", dtype=torch.float32) + try: + newton_schulz_tp( + x_local, + ctx, + num_iterations=5, + partition_dim=1, + tp_mode="distributed", + ) + except ValueError as exc: + if "must be partitioned along their larger dimension" not in str(exc): + raise + else: + raise AssertionError( + "Expected a directly column-sharded tall matrix to be rejected" + ) + if rank == 0: print("Running TP API reference check with replicated input", flush=True) _run_case( diff --git a/transformer_engine/common/newton_schulz/newton_schulz.cpp b/transformer_engine/common/newton_schulz/newton_schulz.cpp index a6af4a1f3bc..740a5e1370f 100644 --- a/transformer_engine/common/newton_schulz/newton_schulz.cpp +++ b/transformer_engine/common/newton_schulz/newton_schulz.cpp @@ -8,6 +8,7 @@ #include +#include #include #include #include @@ -101,6 +102,14 @@ CudaEvent MakeCudaEvent() { } // namespace +struct WorkspaceConfig { + int64_t m; + int64_t n; + cudaDataType_t dtype; + int64_t num_iterations; + std::vector coefficients; +}; + struct NVTECusolverMpCtx { int64_t nranks; int64_t rank; @@ -114,6 +123,8 @@ struct NVTECusolverMpCtx { size_t workspace_size; bool workspace_registered; std::vector workspace_host; + uint64_t* workspace_size_reduction; + std::vector workspace_configs; }; namespace { @@ -141,19 +152,53 @@ size_t GridMaxWorkspaceSize(NVTECusolverMpCtx* ctx, size_t local_size) { } uint64_t size = local_size; - uint64_t* device_size = nullptr; - NVTE_CHECK_CUDA(cudaMalloc(&device_size, sizeof(size))); - NVTE_CHECK_CUDA( - cudaMemcpyAsync(device_size, &size, sizeof(size), cudaMemcpyHostToDevice, ctx->stream.get())); - NVTE_CHECK_NCCL(ncclAllReduce(device_size, device_size, 1, ncclUint64, ncclMax, ctx->comm, - ctx->stream.get())); + if (ctx->workspace_size_reduction == nullptr) { + NVTE_CHECK_CUDA(cudaMalloc(&ctx->workspace_size_reduction, sizeof(size))); + } NVTE_CHECK_CUDA( - cudaMemcpyAsync(&size, device_size, sizeof(size), cudaMemcpyDeviceToHost, ctx->stream.get())); + cudaMemcpyAsync(ctx->workspace_size_reduction, &size, sizeof(size), cudaMemcpyHostToDevice, + ctx->stream.get())); + NVTE_CHECK_NCCL( + ncclAllReduce(ctx->workspace_size_reduction, ctx->workspace_size_reduction, 1, ncclUint64, + ncclMax, ctx->comm, ctx->stream.get())); + NVTE_CHECK_CUDA(cudaMemcpyAsync(&size, ctx->workspace_size_reduction, sizeof(size), + cudaMemcpyDeviceToHost, ctx->stream.get())); NVTE_CHECK_CUDA(cudaStreamSynchronize(ctx->stream.get())); - NVTE_CHECK_CUDA(cudaFree(device_size)); return static_cast(size); } +bool IsWorkspaceConfigCached(const NVTECusolverMpCtx* ctx, int64_t m, int64_t n, + cudaDataType_t dtype, int64_t num_iterations, + const float* coefficients, int64_t num_coefficients) { + const size_t coefficients_size = static_cast(num_coefficients) * sizeof(float); + for (const auto& config : ctx->workspace_configs) { + if (config.m == m && config.n == n && config.dtype == dtype && + config.num_iterations == num_iterations && + config.coefficients.size() == static_cast(num_coefficients) && + (coefficients_size == 0 || + std::memcmp(config.coefficients.data(), coefficients, coefficients_size) == 0)) { + return true; + } + } + return false; +} + +void CacheWorkspaceConfig(NVTECusolverMpCtx* ctx, int64_t m, int64_t n, cudaDataType_t dtype, + int64_t num_iterations, const float* coefficients, + int64_t num_coefficients) { + std::vector cached_coefficients; + if (num_coefficients > 0) { + cached_coefficients.assign(coefficients, coefficients + num_coefficients); + } + ctx->workspace_configs.emplace_back(WorkspaceConfig{ + m, + n, + dtype, + num_iterations, + std::move(cached_coefficients), + }); +} + } // namespace NVTECusolverMpCtx* nvte_cusolvermp_ctx_create(ncclComm_t comm, int nranks, int rank) { @@ -194,12 +239,17 @@ NVTECusolverMpCtx* nvte_cusolvermp_ctx_create(ncclComm_t comm, int nranks, int r 0, false, {}, + nullptr, + {}, }; } void nvte_cusolvermp_ctx_destroy(NVTECusolverMpCtx* ctx) { NVTE_API_CALL(nvte_cusolvermp_ctx_destroy); FreeWorkspace(ctx); + if (ctx->workspace_size_reduction != nullptr) { + NVTE_CHECK_CUDA(cudaFree(ctx->workspace_size_reduction)); + } // Destroy handle and grid before the stream they depend on ctx->grid.reset(); ctx->handle.reset(); @@ -210,6 +260,8 @@ void nvte_newton_schulz(NVTECusolverMpCtx* ctx, int64_t m, int64_t n, NVTETensor int64_t num_iterations, const float* coefficients, int64_t num_coefficients, cudaStream_t caller_stream) { NVTE_API_CALL(nvte_newton_schulz); + NVTE_CHECK(num_iterations >= 0, "Number of iterations must be non-negative, got ", + num_iterations); NVTE_CHECK(num_coefficients == num_iterations * 3, num_iterations, " iterations require ", num_iterations * 3, " coefficients, but ", num_coefficients, " are passed"); const auto* t = convertNVTETensorCheck(x); @@ -246,44 +298,52 @@ void nvte_newton_schulz(NVTECusolverMpCtx* ctx, int64_t m, int64_t n, NVTETensor ns_desc.get(), CUSOLVERMP_NEWTON_SCHULZ_DESCRIPTOR_ATTRIBUTE_REDUCE_VIA_COMPUTE_TYPE, &enabled, sizeof(enabled))); - // Query workspace sizes - size_t wrksp_size_device = 0; - size_t wrksp_size_host = 0; - NVTE_CHECK_CUSOLVERMP(cusolverMpNewtonSchulz_bufferSize( - ctx->handle.get(), ns_desc.get(), n, m, t->data.dptr, 1, 1, mat_desc.get(), num_iterations, - coefficients, CUDA_R_32F, &wrksp_size_device, &wrksp_size_host)); - wrksp_size_device = GridMaxWorkspaceSize(ctx, wrksp_size_device); - - // Allocate/grow device workspace - if (ctx->workspace_size < wrksp_size_device) { - FreeWorkspace(ctx); - - void* workspace = nullptr; - bool workspace_registered = false; - - if (ncclMemAlloc(&workspace, wrksp_size_device) == ncclSuccess) { - if (cusolverMpBufferRegister(ctx->grid.get(), workspace, wrksp_size_device) == - CUSOLVER_STATUS_SUCCESS) { - workspace_registered = true; - } else { - NVTE_CHECK_NCCL(ncclMemFree(workspace)); - workspace = nullptr; + // Workspace requirements are stable for a given operation configuration. Cache configurations + // so repeated optimizer steps avoid a device allocation, collective, and stream synchronization. + const bool workspace_config_cached = + IsWorkspaceConfigCached(ctx, m, n, cuda_dtype, num_iterations, coefficients, + num_coefficients); + if (!workspace_config_cached) { + size_t wrksp_size_device = 0; + size_t wrksp_size_host = 0; + NVTE_CHECK_CUSOLVERMP(cusolverMpNewtonSchulz_bufferSize( + ctx->handle.get(), ns_desc.get(), n, m, t->data.dptr, 1, 1, mat_desc.get(), num_iterations, + coefficients, CUDA_R_32F, &wrksp_size_device, &wrksp_size_host)); + wrksp_size_device = GridMaxWorkspaceSize(ctx, wrksp_size_device); + + // Allocate/grow device workspace + if (ctx->workspace_size < wrksp_size_device) { + FreeWorkspace(ctx); + + void* workspace = nullptr; + bool workspace_registered = false; + + if (ncclMemAlloc(&workspace, wrksp_size_device) == ncclSuccess) { + if (cusolverMpBufferRegister(ctx->grid.get(), workspace, wrksp_size_device) == + CUSOLVER_STATUS_SUCCESS) { + workspace_registered = true; + } else { + NVTE_CHECK_NCCL(ncclMemFree(workspace)); + workspace = nullptr; + } + } + + if (workspace == nullptr) { + NVTE_CHECK_CUDA(cudaMalloc(&workspace, wrksp_size_device)); } - } - if (workspace == nullptr) { - NVTE_CHECK_CUDA(cudaMalloc(&workspace, wrksp_size_device)); + ctx->workspace = workspace; + ctx->workspace_size = wrksp_size_device; + ctx->workspace_registered = workspace_registered; } - ctx->workspace = workspace; - ctx->workspace_size = wrksp_size_device; - ctx->workspace_registered = workspace_registered; - } + // Keep host workspace alive until all work on the internal stream is complete. + if (ctx->workspace_host.size() < wrksp_size_host) { + NVTE_CHECK_CUDA(cudaStreamSynchronize(ctx->stream.get())); + ctx->workspace_host.resize(wrksp_size_host); + } - // Keep host workspace alive until all work on the internal stream is complete. - if (ctx->workspace_host.size() < wrksp_size_host) { - NVTE_CHECK_CUDA(cudaStreamSynchronize(ctx->stream.get())); - ctx->workspace_host.resize(wrksp_size_host); + CacheWorkspaceConfig(ctx, m, n, cuda_dtype, num_iterations, coefficients, num_coefficients); } // Execute Newton-Schulz diff --git a/transformer_engine/pytorch/optimizers/newton_schulz.py b/transformer_engine/pytorch/optimizers/newton_schulz.py index f837627536e..e110d3fffe3 100644 --- a/transformer_engine/pytorch/optimizers/newton_schulz.py +++ b/transformer_engine/pytorch/optimizers/newton_schulz.py @@ -302,7 +302,9 @@ def newton_schulz_tp( tp_mode : {"duplicated", "distributed"}, optional ``"distributed"`` orthogonalizes the existing partition directly. ``"duplicated"`` first gathers the full tensor, orthogonalizes it, and - copies this rank's partition back into ``x``. + copies this rank's partition back into ``x``. In ``"distributed"`` + mode, tall matrices must be partitioned along rows and wide matrices + must be partitioned along columns. """ if x.dim() != 2: raise ValueError(f"Expected 2D tensor, got {x.dim()}D") @@ -331,6 +333,19 @@ def newton_schulz_tp( x.copy_(output) return + if tp_mode == "distributed": + if partition_dim == 0: + global_shape = (x.size(0) * ctx.nranks, x.size(1)) + else: + global_shape = (x.size(0), x.size(1) * ctx.nranks) + if (global_shape[0] > global_shape[1] and partition_dim != 0) or ( + global_shape[0] < global_shape[1] and partition_dim != 1 + ): + raise ValueError( + f"Distributed {global_shape[0]}x{global_shape[1]} matrices must be partitioned " + "along their larger dimension; use tp_mode='duplicated' to redistribute the input" + ) + if tp_mode == "duplicated": x_shards = [torch.empty_like(x) for _ in range(ctx.nranks)] dist.all_gather(x_shards, x, group=ctx.group) From fd8c94e7ff1e46699cd7467058ea8d22d2d8a999 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 18 Sep 2026 20:10:56 +0000 Subject: [PATCH 3/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/distributed/run_newton_schulz.py | 4 +--- .../common/newton_schulz/newton_schulz.cpp | 17 +++++++---------- 2 files changed, 8 insertions(+), 13 deletions(-) diff --git a/tests/pytorch/distributed/run_newton_schulz.py b/tests/pytorch/distributed/run_newton_schulz.py index 2a38bc95f99..22749b5c51a 100644 --- a/tests/pytorch/distributed/run_newton_schulz.py +++ b/tests/pytorch/distributed/run_newton_schulz.py @@ -298,9 +298,7 @@ def run_all_tests(ctx: CusolverMpCtx) -> None: if "must be partitioned along their larger dimension" not in str(exc): raise else: - raise AssertionError( - "Expected a directly column-sharded tall matrix to be rejected" - ) + raise AssertionError("Expected a directly column-sharded tall matrix to be rejected") if rank == 0: print("Running TP API reference check with replicated input", flush=True) diff --git a/transformer_engine/common/newton_schulz/newton_schulz.cpp b/transformer_engine/common/newton_schulz/newton_schulz.cpp index 740a5e1370f..bdc9899803f 100644 --- a/transformer_engine/common/newton_schulz/newton_schulz.cpp +++ b/transformer_engine/common/newton_schulz/newton_schulz.cpp @@ -8,8 +8,8 @@ #include -#include #include +#include #include #include @@ -155,12 +155,10 @@ size_t GridMaxWorkspaceSize(NVTECusolverMpCtx* ctx, size_t local_size) { if (ctx->workspace_size_reduction == nullptr) { NVTE_CHECK_CUDA(cudaMalloc(&ctx->workspace_size_reduction, sizeof(size))); } - NVTE_CHECK_CUDA( - cudaMemcpyAsync(ctx->workspace_size_reduction, &size, sizeof(size), cudaMemcpyHostToDevice, - ctx->stream.get())); - NVTE_CHECK_NCCL( - ncclAllReduce(ctx->workspace_size_reduction, ctx->workspace_size_reduction, 1, ncclUint64, - ncclMax, ctx->comm, ctx->stream.get())); + NVTE_CHECK_CUDA(cudaMemcpyAsync(ctx->workspace_size_reduction, &size, sizeof(size), + cudaMemcpyHostToDevice, ctx->stream.get())); + NVTE_CHECK_NCCL(ncclAllReduce(ctx->workspace_size_reduction, ctx->workspace_size_reduction, 1, + ncclUint64, ncclMax, ctx->comm, ctx->stream.get())); NVTE_CHECK_CUDA(cudaMemcpyAsync(&size, ctx->workspace_size_reduction, sizeof(size), cudaMemcpyDeviceToHost, ctx->stream.get())); NVTE_CHECK_CUDA(cudaStreamSynchronize(ctx->stream.get())); @@ -300,9 +298,8 @@ void nvte_newton_schulz(NVTECusolverMpCtx* ctx, int64_t m, int64_t n, NVTETensor // Workspace requirements are stable for a given operation configuration. Cache configurations // so repeated optimizer steps avoid a device allocation, collective, and stream synchronization. - const bool workspace_config_cached = - IsWorkspaceConfigCached(ctx, m, n, cuda_dtype, num_iterations, coefficients, - num_coefficients); + const bool workspace_config_cached = IsWorkspaceConfigCached( + ctx, m, n, cuda_dtype, num_iterations, coefficients, num_coefficients); if (!workspace_config_cached) { size_t wrksp_size_device = 0; size_t wrksp_size_host = 0;