Sharded data comes back to GPU_0 after pmap

[hyunwoooh5]

Hi. I'm trying to implement data parallelism using jax and I found that all the training set assembles to GPU_0 after pmap.

This is the following code:

import jax
import jax.numpy as jnp
from jax import random, pmap
import optax

# --------------------------------------------------
# Config
num_devices = jax.device_count()
batch_size = 128
per_device_batch = batch_size // num_devices
feature_dim = 16
learning_rate = 1e-2
# --------------------------------------------------

# Dummy Model
def model(params, x):
    return jnp.dot(x, params['w']) + params['b']

# Loss Function
def loss_fn(params, x, y):
    preds = model(params, x)
    return jnp.mean((preds - y)**2)

# Gradient Function (per-device)
@jax.jit
def loss_grad(params, x, y):
    loss, grads = jax.value_and_grad(loss_fn)(params, x, y)
    return grads, loss

# --------------------------------------------------
# Initialization
key = random.PRNGKey(0)
params = {
    'w': random.normal(key, (feature_dim, 1)),
    'b': jnp.zeros((1,))
}

# Replicate params across devices
params_repl = jax.device_put_replicated(params, jax.devices())

# Optax optimizer
opt = optax.sgd(learning_rate)
opt_state = opt.init(params)
opt_state_repl = jax.device_put_replicated(opt_state, jax.devices())

# --------------------------------------------------
# Pre-shard the batch (Example dummy data)
def create_sharded_batch():
    x = random.normal(key, (batch_size, feature_dim))
    y = random.normal(key, (batch_size, 1))

    # Shard manually
    x_shards = jnp.split(x, num_devices)
    y_shards = jnp.split(y, num_devices)

    x_device = jax.device_put_sharded(x_shards, jax.devices())
    y_device = jax.device_put_sharded(y_shards, jax.devices())
    return x_device, y_device

# --------------------------------------------------
# Training Step (pmap version)
def train_step(params, opt_state, x, y):
    grads, loss = loss_grad(params, x, y)
    
    # Average grads across devices
    grads = jax.lax.pmean(grads, axis_name='devices')

    updates, opt_state = opt.update(grads, opt_state)
    params = optax.apply_updates(params, updates)
    return params, opt_state, loss

# Wrap with pmap
parallel_train_step = pmap(train_step, axis_name='devices')

# --------------------------------------------------
# Training Loop
for epoch in range(5):
    x_device, y_device = create_sharded_batch()
    params_repl, opt_state_repl, loss = parallel_train_step(params_repl, opt_state_repl, x_device, y_device)
    print(f"Epoch {epoch}, Loss: {loss[0]}")  # loss is sharded, take first device’s value

Data is sharded when I run create_sharded_batch(), but after training, they are all in the same basket.

Question is

1. Is it the property of pmap?
2. Is there anyway I can avoid this to maximize GPU performance by the training set not moving among GPUs?

[yashk2810]

I would suggest using jax.jit and if you want to write your own collectives, then use shard_map

Some docs for how to do that:

https://docs.jax.dev/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html

https://docs.jax.dev/en/latest/notebooks/explicit-sharding.html

https://docs.jax.dev/en/latest/notebooks/shard_map.html