MPI Guide

This page explains how to run JAX-AMG in distributed mode with MPI across multiple GPUs and nodes.

If your application already uses JAX sharding, see the additive JAX Sharding interface. It retains MPI inside AmgX rather than replacing the workflow documented here.

Prerequisites

Before running MPI jobs, make sure you have:

A working MPI installation (mpirun available).
mpi4py and mpi4jax installed in the same environment as jaxamg.
CUDA-enabled JAX and a working NVIDIA driver/toolkit.
AmgX installed and visible through LD_LIBRARY_PATH.

See the Installation Guide for more details.

Launching MPI Jobs

To run an MPI job, use:

mpirun -n 2 python demo/mpi_tridiagonal_matrix_optimization.py

You can set CUDA_VISIBLE_DEVICES to control which GPUs are visible to each process.

Distributed Solve Pattern

JAX-AMG MPI solves follow this pattern:

Initialize communicator (MPI.COMM_WORLD).
Build a local matrix partition per rank.
Build local right-hand-side vector.
Call jaxamg.solve(...) with comm, nglobal, and partition_info.
Gather local solutions if you need a global vector on rank 0.

Example:

from mpi4py import MPI
import jaxamg
from jaxamg.matrices import poisson_matrix_distributed, rhs_ones
from jaxamg.mpi_utils import partition_vector, gather_vector

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
nranks = comm.Get_size()

# Build local matrix rows for each rank
n = 16
A_local, row_start, row_end = poisson_matrix_distributed(n, n, rank, nranks)
b_local, _, _ = partition_vector(rhs_ones(n * n), rank, nranks)

config = {
    "solver": "CG",
    "preconditioner": {"solver": "JACOBI_L1"},
    "communicator": "MPI_DIRECT",
}

# Solve
x_local, info = jaxamg.solve(
    A_local,
    b_local,
    comm=comm,
    nglobal=n * n,
    partition_info=(row_start, row_end),
    config=config,
)

# Gather solution
x_global = gather_vector(x_local, comm, root=0)
if rank == 0: print(f"Soluton: {x_global}")

# Finalization
comm.Barrier()
jaxamg.finalize()

Block matrices (block_dim > 1) work in MPI mode as well; each rank's local partition must be block-aligned (its row count divisible by block_dim). Note that the AMG coarse-solver default differs from the single-GPU case — see Block matrices in the configuration guide.

Caching

If you are solving many similar systems repeatedly, such as in an optimization loop, especially with JIT, compute and cache the metadata once (outside the JIT-compiled region) and reuse it for subsequent solves:

# Compute once (outside JIT)
mpi_cache = jaxamg.cache_mpi_metadata(
    solver_config, comm, n_global, partition_info, A_local
)

# Use inside JIT if needed
A_cached = jaxamg.with_cache(A_local, mpi=mpi_cache, is_symmetric=True)
x, info = jaxamg.solve(A_cached, b_local)

In addition, you can tune the native resource cache behavior (which is separate from the metadata caching above) using:

export JAXAMG_CACHE_SIZE=2

Set 0 to disable resource caching, or use larger values to keep more solver resources alive between calls.

See the Caching Guide for more details.

GPU-Aware MPI

To enable GPU-aware MPI, first verify that your MPI library was compiled with CUDA support. For OpenMPI, you can check with:

ompi_info --parsable --all | grep mpi_built_with_cuda_support:value

Next, set the necessary environment variables for GPU-aware operation:

export OMPI_MCA_opal_cuda_support=true   # required for OpenMPI
export MPI4JAX_USE_CUDA_MPI=1            # for mpi4jax

When using AmgX with GPU-aware MPI, specify the communicator in the configuration:

config = {
    "communicator": "MPI_DIRECT",
    ...
}

Non-GPU-aware MPI still works, but communication may stage through host memory and be slower.

Singular systems

nullspace and transpose_nullspace (see Singular systems) take each rank's local rows; the projections reduce across ranks internally. With cached metadata, prepare it with singular=True and attach the bases to A:

mpi_cache = jaxamg.cache_mpi_metadata(config, comm, n, (row_start, row_end), A_local, singular=True)
A_local = jaxamg.with_cache(A_local, mpi=mpi_cache, nullspace="constant", transpose_nullspace=V_local)
x_local, info = jaxamg.solve(A_local, b_local)

Labels contain each rank's local rows. Use the same label for a part on every rank; totals are reduced globally. For custom numbering, pass label_sum, a linear JAX function mapping local per-label sums to totals. Use jaxamg.nullspace.make_mpi_reduce_sum for differentiable MPI sums. With custom numbering, cross-label coupling checks cover only locally owned columns.

Requirements and limits

Apply the same differentiation transforms on every rank. Keep zero dependencies where needed (e.g. b + 0 * theta and loss + 0 * x.sum()) so no rank skips the solve's collectives.
Clear solver caches on all ranks together.
Use one host thread per communicator. Keep the communicator alive until queued work finishes and all ranks call jaxamg.finalize().
Restart the process after a native-call failure.
Local scalar row and nonzero counts must fit signed 32-bit integers. Published AmgX releases also limit the global block-row count (nglobal / block_dim) to signed 32-bit. To exceed that limit, build the int64-global-rows AMGX port and rebuild JAX-AMG against it. This port's AMGX_matrix_upload_distributed accepts a 64-bit n_global; use the matching library at runtime. The header selects the upload width automatically; stock builds retain their 32-bit upload path and reject larger systems. Global columns remain 64-bit on the Python/FFI boundary in either case. This does not enable larger local partitions or single-GPU systems. Each exchange message must be smaller than 2 GiB.

Halo-form operators

Use halo_operator for apply(x_local, x_ghost). Supply sorted global column IDs in ghost_ids; JAX-AMG fetches those entries. apply must not communicate.

op = jaxamg.halo_operator(apply, n_local=n_local, ghost_ids=ghost_ids)
x, info = jaxamg.solve(op, b_local, comm=comm, nglobal=n, partition_info=(r0, r1))

Before jit or grad, solve once eagerly or declare pattern=. Reuse the structure with op.with_fn(new_apply). See demo/mpi_halo_operator_optimization.py.