Train with JAX on Multiple Heterogeneous GPUs (RTX 4060Ti and RTX 5060Ti)

[animikhaich]

I have two devices:

• GPU 0: RTX 4060Ti 16 GB
• GPU 1: RTX 5060Ti 16 GB

With a background in PyTorch, I have used PyTorch to train on both GPUs together using DDP.

I was playing around with Jax and Flax, and hence started with the official MNIST Example on Flax. It ran fine on either GPUs individually, and wanted test out distributed training across both my GPUs.

With Google Gemini's help, I ended up with this code:

# Suppress warnings for a cleaner output
import os
import warnings
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'  # Suppress TensorFlow C++ logs
warnings.filterwarnings("ignore")  # Suppress standard Python warnings

import tensorflow_datasets as tfds
import tensorflow as tf
from tqdm import tqdm
import jax
import jax.numpy as jnp
import optax
from flax import nnx
from functools import partial
import matplotlib.pyplot as plt

# --- Multi-GPU Setup ---
# Prevent TensorFlow from occupying GPU memory, as JAX will be managing it.
tf.config.experimental.set_visible_devices([], 'GPU')
# Get the number of available JAX devices (GPUs/TPUs)
num_devices = jax.local_device_count()
print(f"Running on {num_devices} devices.")

# --- Hyperparameters ---
train_steps = 12000
eval_every = 200
# The total batch size will be split across all available devices.
# Ensure the total batch size is divisible by the number of devices.
total_batch_size = 256 * num_devices
learning_rate = 0.005
momentum = 0.9

# --- Data Loading and Preprocessing ---
def preprocess(sample):
  """Normalize images."""
  return {
      'image': tf.cast(sample['image'], tf.float32) / 255.0,
      'label': sample['label'],
  }

train_ds: tf.data.Dataset = tfds.load('mnist', split='train').map(preprocess)
test_ds: tf.data.Dataset = tfds.load('mnist', split='test').map(preprocess)

train_ds = train_ds.repeat().shuffle(1024)
train_ds = train_ds.batch(total_batch_size, drop_remainder=True).take(train_steps).prefetch(tf.data.AUTOTUNE)
test_ds = test_ds.batch(total_batch_size, drop_remainder=True).prefetch(tf.data.AUTOTUNE)

def shard(batch):
  """Shards a batch of data across all available devices."""
  return jax.tree.map(
      lambda x: x.reshape((num_devices, -1) + x.shape[1:]), batch
  )

# --- Model Definition (Unchanged) ---
class CNN(nnx.Module):
  """A simple CNN model."""
  def __init__(self, *, rngs: nnx.Rngs):
    self.conv1 = nnx.Conv(1, 32, kernel_size=(3, 3), rngs=rngs)
    self.batch_norm1 = nnx.BatchNorm(32, use_running_average=False, rngs=rngs)
    self.dropout1 = nnx.Dropout(rate=0.025, rngs=rngs)
    self.conv2 = nnx.Conv(32, 64, kernel_size=(3, 3), rngs=rngs)
    self.batch_norm2 = nnx.BatchNorm(64, use_running_average=False, rngs=rngs)
    self.avg_pool = partial(nnx.avg_pool, window_shape=(2, 2), strides=(2, 2))
    self.linear1 = nnx.Linear(3136, 256, rngs=rngs)
    self.dropout2 = nnx.Dropout(rate=0.025, rngs=rngs)
    self.linear2 = nnx.Linear(256, 10, rngs=rngs)

  def __call__(self, x, *, train: bool):
    # Pass the `train` flag to BatchNorm and Dropout layers
    self.batch_norm1.use_running_average = not train
    self.batch_norm2.use_running_average = not train
    self.dropout1.deterministic = not train
    self.dropout2.deterministic = not train

    x = self.avg_pool(nnx.relu(self.batch_norm1(self.dropout1(self.conv1(x)))))
    x = self.avg_pool(nnx.relu(self.batch_norm2(self.conv2(x))))
    x = x.reshape(x.shape[0], -1)  # flatten
    x = nnx.relu(self.dropout2(self.linear1(x)))
    x = self.linear2(x)
    return x

# --- State Management ---
class TrainState(nnx.Object):
  """A container for the complete training state."""
  def __init__(self, model: CNN, optimizer: nnx.Optimizer, metrics: nnx.MultiMetric):
    self.model = model
    self.optimizer = optimizer
    self.metrics = metrics

# --- Loss Function ---
def loss_fn(model: CNN, batch):
  logits = model(batch['image'], train=True)
  loss = optax.softmax_cross_entropy_with_integer_labels(
      logits=logits, labels=batch['label']
  ).mean()
  return loss, logits

# --- Training and Evaluation Steps (Parallelized) ---
@partial(nnx.pmap, axis_name='batch')
def train_step(state: TrainState, batch):
  """Train for a single step, parallelized across devices."""
  grad_fn = nnx.value_and_grad(loss_fn, has_aux=True)
  (loss, logits), grads = grad_fn(state.model, batch)

  # Aggregate gradients and loss from all devices
  grads = jax.lax.pmean(grads, axis_name='batch')
  loss = jax.lax.pmean(loss, axis_name='batch')

  # Update metrics and optimizer (in-place)
  state.metrics.update(loss=loss, logits=logits, labels=batch['label'])
  state.optimizer.update(state.model, grads)
  return state

@partial(nnx.pmap, axis_name='batch')
def eval_step(state: TrainState, batch):
  """Evaluate for a single step, parallelized across devices."""
  logits = state.model(batch['image'], train=False)
  loss = optax.softmax_cross_entropy_with_integer_labels(
      logits=logits, labels=batch['label']
  ).mean()

  # Aggregate loss from all devices
  loss = jax.lax.pmean(loss, axis_name='batch')

  # Update metrics (in-place)
  state.metrics.update(loss=loss, logits=logits, labels=batch['label'])
  return state

# --- Initialization ---
# Initialize model, optimizer, and metrics on the CPU first
model = CNN(rngs=nnx.Rngs(0))
optimizer = nnx.Optimizer(model, optax.adamw(learning_rate, momentum), wrt=nnx.Param)
metrics = nnx.MultiMetric(
    accuracy=nnx.metrics.Accuracy(),
    loss=nnx.metrics.Average('loss'),
)

# Create and replicate the training state across all devices
state = TrainState(model=model, optimizer=optimizer, metrics=metrics)
replicated_state = jax.device_put_replicated(state, jax.local_devices())

# --- Training Loop ---
metrics_history = {
    'train_loss': [], 'train_accuracy': [],
    'test_loss': [], 'test_accuracy': [],
}

print("Starting training...")
for step, batch in tqdm(enumerate(train_ds.as_numpy_iterator()), total=train_steps):
  batch = shard(batch)  # Distribute the batch to all devices
  replicated_state = train_step(replicated_state, batch)

  if (step + 1) % eval_every == 0 or (step + 1) == train_steps:
    # Get metrics from one device (they are all synchronized)
    current_metrics = jax.device_get(replicated_state.metrics)
    for metric, value in current_metrics.compute().items():
      metrics_history[f'train_{metric}'].append(value)
    current_metrics.reset()

    # Evaluation loop
    for test_batch in test_ds.as_numpy_iterator():
      test_batch = shard(test_batch)
      replicated_state = eval_step(replicated_state, test_batch)

    current_metrics = jax.device_get(replicated_state.metrics)
    for metric, value in current_metrics.compute().items():
      metrics_history[f'test_{metric}'].append(value)
    current_metrics.reset()

print("Training finished.")
# --- Plotting Results ---
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))
ax1.set_title('Loss')
ax2.set_title('Accuracy')
for dataset in ('train', 'test'):
  ax1.plot(metrics_history[f'{dataset}_loss'], label=f'{dataset}_loss')
  ax2.plot(metrics_history[f'{dataset}_accuracy'], label=f'{dataset}_accuracy')
ax1.legend()
ax2.legend()
plt.show()

I got the following error:

WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
W0000 00:00:1758609847.735368   65509 gpu_device.cc:2431] TensorFlow was not built with CUDA kernel binaries compatible with compute capability 12.0. CUDA kernels will be jit-compiled from PTX, which could take 30 minutes or longer.
Running on 2 devices.
Starting training...
  0%|                                                                                                                        | 0/12000 [00:03<?, ?it/s]
Traceback (most recent call last):
  File "/home/ani/Projects/videoprism/jax_test.py", line 151, in <module>
    replicated_state = train_step(replicated_state, batch)
  File "/home/ani/Projects/videoprism/.venv/lib/python3.13/site-packages/flax/nnx/graph.py", line 2045, in update_context_manager_wrapper
    return f(*args, **kwargs)
  File "/home/ani/Projects/videoprism/.venv/lib/python3.13/site-packages/flax/nnx/transforms/iteration.py", line 573, in vmap_wrapper
    pure_args_out, pure_out = pmapped_fn(*pure_args)
                              ~~~~~~~~~~^^^^^^^^^^^^
jaxlib._jax.XlaRuntimeError: INVALID_ARGUMENT: executable is built for device CUDA:0 of type "NVIDIA GeForce RTX 4060 Ti"; cannot run it on device CUDA:1 of type "NVIDIA GeForce RTX 5060 Ti"
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

Did some digging and realized that Data Parallelization is not supported across heterogenous hardware with Jax. Is that true? Is there any way around it?