Distributed Linear Algebra#

The distributed Linear Algebra APIs in nvmath-python leverage NVIDIA’s high-performance distributed math libraries to perform dense linear algebra on multi-node multi-GPU systems at scale, directly from host code. They currently provide:

Matrix multiplication, built on the NVIDIA cuBLASMp library.
Direct linear solver, built on the NVIDIA cuSOLVERMp library.

Each operation offers both stateless function-form APIs and stateful class-form APIs. As with all distributed APIs, the operands on each process are the local partition of the global operands, and the user specifies the distribution.

Operand distribution. These APIs natively use the block-cyclic distribution (see Block distributions), so the distribution you specify for each operand must be compatible with it. The supported choices are BlockCyclic, BlockNonCyclic, and Slab (with uniform partition sizes).

Memory layout. cuBLASMp and cuSOLVERMp require operands to use Fortran-order memory layout, while Python libraries such as NumPy and PyTorch use C-order by default. See Distribution, memory layout and transpose for guidelines on memory layout conversion for distributed operands and potential implications on distribution.

Matrix Multiplication#

The distributed matrix multiplication APIs (in nvmath.distributed.linalg.advanced) leverage the NVIDIA cuBLASMp library to compute \(F(\alpha A @ B + \beta C)\), where \(F\) is an optional epilog, on multi-node multi-GPU systems, with both a function-form (matmul()) and a stateful class-form (Matmul) API.

Matrix qualifiers#

Matrix qualifiers are used to indicate whether an input matrix is transposed or not.

For example, for A.T @ B you have to specify:

import numpy as np
from nvmath.distributed.linalg.advanced import matrix_qualifiers_dtype, matmul

qualifiers = np.zeros((3,), dtype=matrix_qualifiers_dtype)
qualifiers[0]["is_transpose"] = True  # a is transposed
qualifiers[1]["is_transpose"] = False  # b is not transposed (optional)

...

result = matmul(a, b, distributions=distributions, qualifiers=qualifiers)

Caution

A common strategy to convert memory layout to Fortran-order (required by cuBLASMp) is to transpose the input matrices, as explained in Distribution, memory layout and transpose. Remember to set the matrix qualifiers accordingly.

Distributed algorithm#

cuBLASMp implements efficient communication-overlap algorithms that are suited for distributed machine learning scenarios with tensor parallelism. Algorithms include AllGather+GEMM and GEMM+ReduceScatter. These algorithms have special requirements in terms of how each of the operands is distributed and their transpose qualifiers.

Currently, to be able to use these algorithms, cuBLASMp requires: the matrices to be distributed using a 1D partitioning scheme without the cyclic distribution and the partition sizes to be uniform (BlockNonCyclic and Slab are valid distributions for this use case).

Please refer to cuBLASMp documentation for full details.

Epilogue input and output distribution#

Generally, the distribution of an epilogue input or output is the same as the distribution of the matrix associated with that input/output. For example, the bias vector is applied to the matmul result and has the same distribution: if the result matrix is partitioned on the M dimension, the bias is similarly partitioned; if the result is partitioned on the N dimension (but not M) the bias is replicated on every process.

For bias gradient epilogues, the epilogue output distribution is as follows:

For BGRAD and BGRADA it follows how matrix A is partitioned on M.
For BGRADB the epilogue output is always replicated.

Example#

The following example performs \(\alpha A @ B + \beta C\) with inputs distributed according to a Slab distribution (partitioning on a single dimension):

Note

To use the distributed Matmul APIs you need to initialize the distributed runtime with the NCCL communication backend.

import numpy as np
import cupy as cp
from nvmath.distributed.distribution import Slab
from nvmath.distributed.linalg.advanced import matrix_qualifiers_dtype

# Get my process rank.
rank = nvmath.distributed.get_context().process_group.rank

# The global problem size m, n, k
m, n, k = 128, 512, 1024

# Prepare sample input data.
with cp.cuda.Device(device_id):
    a = cp.random.rand(*Slab.X.shape(rank, (m, k)))
    b = cp.random.rand(*Slab.X.shape(rank, (n, k)))
    c = cp.random.rand(*Slab.Y.shape(rank, (n, m)))

# Get transposed views with Fortran-order memory layout
a = a.T  # a is now (k, m) with Slab.Y
b = b.T  # b is now (k, n) with Slab.Y
c = c.T  # c is now (m, n) with Slab.X

distributions = [Slab.Y, Slab.Y, Slab.X]

qualifiers = np.zeros((3,), dtype=matrix_qualifiers_dtype)
qualifiers[0]["is_transpose"] = True  # a is transposed

alpha = 0.45
beta = 0.67

# Perform the distributed GEMM.
result = nvmath.distributed.linalg.advanced.matmul(
    a,
    b,
    c=c,
    alpha=alpha,
    beta=beta,
    distributions=distributions,
    qualifiers=qualifiers,
)

# Distributed matrix multiplication is inplace by default when c
# is provided (the result is stored in c). This behavior can
# be modified using the `inplace` option.
assert result is c

# Synchronize the default stream, since by default the execution
# is non-blocking for GPU operands.
cp.cuda.get_current_stream().synchronize()

# result is distributed row-wise
assert result.shape == Slab.X.shape(rank, (m, n))

You can find many more examples here.

API Reference#

`matmul`(a, b, /[, c, alpha, beta, epilog, ...])	Perform the specified distributed matrix multiplication computation \(\mathcal{F}(\alpha a @ b + \beta c)\), where \(\mathcal{F}\) is the epilog.
`matrix_qualifiers_dtype`	NumPy dtype object that encapsulates the matrix qualifiers in distributed.linalg.advanced.
`Matmul`(a, b, /[, c, alpha, beta, ...])	Create a stateful object encapsulating the specified distributed matrix multiplication computation \(\alpha a @ b + \beta c\) and the required resources to perform the operation.
`MatmulComputeType`	alias of `ComputeType`
`MatmulEpilog`	alias of `MatmulEpilogue`
`MatmulAlgoType`(*values)	See `cublasMpMatmulAlgoType_t`.

`MatmulEpilogPreferences`([aux_type, aux_amax])	A data class for providing epilog options as part of `preferences` to the `Matmul.plan()` method and the wrapper function `matmul()`.
`MatmulOptions`([inplace, compute_type, ...])	A data class for providing options to the `Matmul` object and the wrapper function `matmul()`.
`MatmulPlanPreferences`([epilog])	A data class for providing options to the `Matmul.plan()` method and the wrapper function `matmul()`.
`MatmulQuantizationScales`([a, b, c, d])	A data class for providing quantization_scales to `Matmul` constructor and the wrapper function `matmul()`.

Helpers#

The module nvmath.linalg.advanced.helpers.matmul provides helper functions that facilitate working with the narrow-precision features of nvmath.distributed.linalg.advanced module.

Direct Linear Solver#

The distributed direct linear solver APIs (in nvmath.distributed.linalg) leverage the NVIDIA cuSOLVERMp library to solve dense linear systems \(A @ X = B\) on multi-node multi-GPU systems, via an LU factorization with partial pivoting followed by triangular solves. Both a function-form (direct_solver()) and a stateful class-form (DirectSolver) API are provided. On each process, the operands a and b hold only the local partitions of \(A\) and \(B\) according to the specified distributions; see DirectSolver for the supported distributions and the constraints they must satisfy.

Note

To use the distributed direct solver APIs you need to initialize the distributed runtime with the NCCL communication backend.

You can find examples here.

API Reference#

`direct_solver`(a, b, /, *, distributions[, ...])	Solve the distributed dense linear system \(a @ x = b\) for \(x\).
`DirectSolver`(a, b, /, *, distributions[, ...])	Create a stateful object that encapsulates the specified distributed direct solver computation for the dense linear system \(a @ x = b\) and the required resources to perform it.
`InvalidDirectSolverState`	Raised when a public method is called after `DirectSolver.free()`.

DirectSolverOptions(*, logger, blocking, ...)

A data class for providing options to direct_solver() and DirectSolver.