Question on Load Imbalance Using JAX with MPI

[kousuke-nakano]

Dear users and developers,

I'm currently performing parallel computations using JAX combined with MPI (via mpi4py) on a CPU cluster machine. I'm encountering issues related to load imbalance among MPI processes, which is puzzling given that each process executes exactly the same computation.

System Info:

• Compute node: CPU Intel Xeon Platinum 8490H (Sapphire Rapids, 60 cores, 1.90 GHz - 3.50 GHz) × 2 sockets
• Memory: DDR5 4800 MHz, 512 GiB (32GiB×8 ×2 sockets)
• Memory bandwidth: 614.2 GB/s （4800MHz×8Byte×8 channels×2 sockets）
• Interconnect: InfiniBand NDR (200Gbps) × 1 port/node

Python Info:

%python --version
Python 3.11.0

%pip freeze
jax==0.5.3
jaxlib==0.5.3
ml_dtypes==0.5.1
mpi4py==4.0.3
numpy==2.2.4
opt_einsum==3.4.0
scipy==1.15.2

Below is a simplified version of my JAX-MPI Python program. Since MPI communications are needed outside jitted functions in my code, I use mpi4py, not mpi4jax.

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import time
import jax
import jax.numpy as jnp
from jax import jit, vmap
import numpy as np
from mpi4py import MPI

# Initialize MPI
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

# jax-MPI related initialization
jax.distributed.initialize(cluster_detection_method="mpi4py")

# Gather JAX device information
global_device_info = jax.devices()
local_device_info = jax.local_devices()

'''
# Only rank 0 prints the recognized XLA devices
if rank == 0:
    print("*** XLA Global devices recognized by JAX ***")
    print(global_device_info)
    print("*** XLA Local devices recognized by JAX ***")
    print(local_device_info)
    print("")
'''

# Insert a barrier after device information printing
comm.Barrier()

# JIT-wrapped function using a for-loop to perform matrix multiplications n_iter times
@jit
def compute_matrix_operations(a, b, n_iter=10):
    for _ in range(n_iter):
        a = jnp.dot(a, b)
    return a

# Set the number of copies to process simultaneously
N = 100
M = 100

# Generate a batch of matrices for each MPI process (shape: (N, M, M))
key = jax.random.PRNGKey(1234)
key_a, key_b = jax.random.split(key)
a_batch = jax.random.uniform(key_a, (N, M, M))
b_batch = jax.random.uniform(key_b, (N, M, M))

# Create a vectorized version of compute_matrix_operations using vmap
vmap_compute = vmap(compute_matrix_operations, in_axes=(0, 0))

# Measure compilation time (includes the JIT compilation overhead)
start_compile = time.perf_counter()
compiled_result = vmap_compute(a_batch, b_batch)
compiled_result = compiled_result.block_until_ready()
compile_time = time.perf_counter() - start_compile

# Measure net JAX computation time (after JIT has compiled) in seconds
start_compute = time.perf_counter()
result = vmap_compute(a_batch, b_batch)
result = result.block_until_ready()  # Wait until the computation is complete
compute_time = time.perf_counter() - start_compute

# Measure MPI Barrier time in seconds
start_barrier = time.perf_counter()
comm.Barrier()
barrier_time = time.perf_counter() - start_barrier

# Measure MPI Allreduce time in seconds
result_np = np.array(result)  # Convert JAX result to a numpy array for MPI communication
total_result = np.empty_like(result_np)
start_allreduce = time.perf_counter()
comm.Allreduce(result_np, total_result, op=MPI.SUM)
allreduce_time = time.perf_counter() - start_allreduce

# Gather the minimum and maximum times across all processes
compile_time_min = comm.allreduce(compile_time, op=MPI.MIN)
compile_time_max = comm.allreduce(compile_time, op=MPI.MAX)
compute_time_min = comm.allreduce(compute_time, op=MPI.MIN)
compute_time_max = comm.allreduce(compute_time, op=MPI.MAX)
barrier_time_min = comm.allreduce(barrier_time, op=MPI.MIN)
barrier_time_max = comm.allreduce(barrier_time, op=MPI.MAX)
allreduce_time_min = comm.allreduce(allreduce_time, op=MPI.MIN)
allreduce_time_max = comm.allreduce(allreduce_time, op=MPI.MAX)

# Only rank 0 prints the aggregated timing information in milliseconds
if rank == 0:
    print("Compilation time:      min = {:.6f} msec, max = {:.6f} msec".format(
        compile_time_min * 1000, compile_time_max * 1000))
    print("JAX computation time:  min = {:.6f} msec, max = {:.6f} msec".format(
        compute_time_min * 1000, compute_time_max * 1000))
    print("MPI Barrier time:      min = {:.6f} msec, max = {:.6f} msec".format(
        barrier_time_min * 1000, barrier_time_max * 1000))
    print("MPI Allreduce time:    min = {:.6f} msec, max = {:.6f} msec".format(
        allreduce_time_min * 1000, allreduce_time_max * 1000))

When running this script on 32 nodes (3840 MPI processes),

#!/bin/sh
#PJM --mpi proc=3840
#PJM -L node=32
#PJM -L elapse=00:01:00
...
# load intel mpi
module purge
module load intel/2023.2
module load impi/2021.10.0
# env
export OMP_NUM_THREADS=1
export MKL_NUM_THREADS=1
export XLA_FLAGS="--xla_force_host_platform_device_count=1 --xla_cpu_multi_thread_eigen=false intra_op_parallelism_threads=1"
# run
mpiexec --bind-to core -np 3840 python jax_mpi.py > out_on_cpu_node_32 2> out_on_cpu_node_32.e

where I bound each MPI process to each core (by —bind-to-core) and switched off the multithreading (by export XLA_FLAGS="--xla_force_host_platform_device_count=1 --xla_cpu_multi_thread_eigen=false intra_op_parallelism_threads=1”).

I'm observing significant load imbalance.

Compilation time:      min = 298.444553 msec, max = 681.890790 msec
JAX computation time:  min = 30.428662 msec, max = 78.641703 msec
MPI Barrier time:      min = 0.104042 msec, max = 704.886884 msec
MPI Allreduce time:    min = 21.611117 msec, max = 38.786539 msec

Given the theoretical expectation that all processes should be completed simultaneously before the MPI barrier (since they execute identical workloads), the high barrier time indicates severe load imbalance.

Questions:

1. Is my timing measurement method (inserting .block_until_ready()) correct to accurately measure JAX computation times?
2. Could this load imbalance be related to JAX’s native multi-threading behavior or some internal JAX/XLA scheduling causing uneven execution times, even though I switched them off?
3. Do you have any suggestions or tools that could help identify and resolve the cause of this load imbalance?

Any advice or suggestions would be greatly appreciated!