# How to Use Distributed Training with Multiple GPUs and Data Sharding in EvoRL

> Learn distributed training with multiple GPUs and data sharding in EvoRL. Optimize large-scale population evolution using JAX pmap and shard_map primitives. Get started today!

- Repository: [EMI-Group/evorl](https://github.com/emi-group/evorl)
- Tags: how-to-guide
- Published: 2026-03-01

---

**EvoRL enables distributed training across multiple GPUs by automatically replicating workflow states and averaging gradients via JAX's `pmap` primitives, while data sharding is handled through `shard_map` wrappers for large-scale population evolution.**

EvoRL is an open-source evolutionary reinforcement learning framework built on JAX that abstracts away the complexity of multi-device execution. Whether you are running population-based evolution or standard RL algorithms, the framework provides a seamless path from single-GPU prototyping to multi-GPU scaling without changing your algorithm implementation.

## How EvoRL Handles Multi-GPU Distribution

The distributed architecture centers on the **RLWorkflow** base class in [`evorl/workflows/rl_workflow.py`](https://github.com/emi-group/evorl/blob/main/evorl/workflows/rl_workflow.py). This class inspects the runtime environment and adapts the training loop accordingly.

When `enable_multi_devices` is set to `True`, the workflow transforms its `step` and `evaluate` methods using `jax.pmap` over the axis defined by **PMAP_AXIS_NAME** (defined in [`evorl/distributed/__init__.py`](https://github.com/emi-group/evorl/blob/main/evorl/distributed/__init__.py)). This global axis name ensures all cross-device reductions use a consistent identifier.

The communication primitives live in [`evorl/distributed/comm.py`](https://github.com/emi-group/evorl/blob/main/evorl/distributed/comm.py):

- **`split_key_to_devices`**: Splits a JAX PRNG key into device-specific sub-keys, ensuring statistical independence across replicas.
- **`pmean`, `psum`, `pmax`, `pmin`**: Thin wrappers around `jax.lax.p*` that become no-ops on single-device runs but perform collective operations when under `pmap`.
- **`gradient_update`**: Located in [`evorl/distributed/gradients.py`](https://github.com/emi-group/evorl/blob/main/evorl/distributed/gradients.py), this utility computes local gradients and automatically averages them across devices using `pmean`, keeping parameters synchronized.

## Setting Up Distributed Training

### Automatic Device Detection

The simplest entry point is [`scripts/train.py`](https://github.com/emi-group/evorl/blob/main/scripts/train.py), which automatically enables multi-device mode when JAX detects more than one local GPU.

```bash
python -m scripts.train \
    +env=cartpole \
    workflow_cls=evorl.workflows.rl_workflow.OnPolicyWorkflow \
    enable_jit=true

```

When multiple devices are present, the script prints a confirmation:

```

Enable Multiple Devices: [GpuDevice(id=0), GpuDevice(id=1), GpuDevice(id=2), GpuDevice(id=3)]

```

### Hydra Configuration Override

For explicit control, create a configuration overlay at [`configs/multi_gpu.yaml`](https://github.com/emi-group/evorl/blob/main/configs/multi_gpu.yaml):

```yaml
defaults:
  - _self_
  - config

enable_multi_devices: true

```

Launch with:

```bash
python -m scripts.train +multi_gpu.yaml

```

### State Replication Internals

Inside `OnPolicyWorkflow.setup` or `OffPolicyWorkflow.setup`, the workflow executes the following when `self.enable_multi_devices` is active:

1. Replicates agent state, optimizer state, and metrics via `jax.device_put_replicated`.
2. Distributes unique PRNG keys using `split_key_to_devices`.
3. Initializes environment states independently on each device via `pmap`.

This ensures every GPU maintains an identical copy of the model while processing different environment instances.

## Data Sharding with shard_map

For algorithms requiring explicit data partitioning—such as population-based evolution with large genotype arrays—EvoRL provides `shmap_vmap` and `shmap_map` in [`evorl/distributed/sharding.py`](https://github.com/emi-group/evorl/blob/main/evorl/distributed/sharding.py). These wrap `jax.experimental.shard_map` to distribute data across the device mesh.

```python
import jax
import jax.numpy as jnp
from evorl.distributed import shmap_vmap, tree_device_put

# Create a mesh spanning all available GPUs

mesh = jax.sharding.Mesh(jax.devices(), ('gpu',))

def evaluate_population(pop_params):
    # Dummy fitness evaluation

    return jnp.sum(pop_params, axis=-1)

# Define sharding specs: distribute the first axis across GPUs

in_specs = (jax.sharding.PartitionSpec('gpu',),)
out_specs = (jax.sharding.PartitionSpec('gpu',),)

# Wrap the function

sharded_eval = shmap_vmap(evaluate_population, mesh, in_specs, out_specs)

# Prepare data: [num_devices, pop_per_device, param_dim]

pop = jnp.arange(4 * 10 * 3).reshape(4, 10, 3).astype(jnp.float32)
pop_sharded = tree_device_put(pop, mesh)

# Execute: each GPU processes its shard in parallel

fitness = sharded_eval(pop_sharded)  # Shape: (4, 10)

```

## Gradient Synchronization Across Devices

During the training step, gradients computed on each device must be averaged to maintain model consistency. EvoRL handles this transparently through the `pmean` utility.

```python
from evorl.distributed import pmean
import optax

def train_step(state, batch):
    # Compute local gradients

    grads = jax.grad(loss_fn)(state.params, batch)
    
    # Average across all devices using the PMAP_AXIS_NAME

    grads = pmean(grads, axis_name="P")
    
    # Apply optimizer update

    updates, new_opt_state = state.optimizer.update(grads, state.opt_state)
    new_params = optax.apply_updates(state.params, updates)
    
    return state.replace(params=new_params, opt_state=new_opt_state)

```

When `train_step` is transformed by `pmap`, the `pmean` call performs an all-reduce operation. On single-device runs, the call incurs zero overhead.

## Multi-Node Cluster Support

For multi-node clusters where each process controls one GPU, use [`scripts/train_dist.py`](https://github.com/emi-group/evorl/blob/main/scripts/train_dist.py). This entry point assumes external launcher coordination (such as `torchrun` or `jaxrun`) and queries `is_dist_initialized()` from [`evorl/distributed/comm.py`](https://github.com/emi-group/evorl/blob/main/evorl/distributed/comm.py) to detect global ranks.

```python
from evorl.distributed import get_global_ranks

ranks = get_global_ranks()  # Returns array of ranks across all nodes

local_rank = ranks[0]       # Each process extracts its identifier

# Use local_rank to partition dataset shards or population slices

```

## Summary

- **RLWorkflow** automatically enables multi-GPU execution via the `enable_multi_devices` flag, transforming methods with `pmap` without code changes.
- **Communication primitives** (`pmean`, `psum`, `split_key_to_devices`) in [`evorl/distributed/comm.py`](https://github.com/emi-group/evorl/blob/main/evorl/distributed/comm.py) handle cross-device reduction and key distribution.
- **State replication** occurs automatically in workflow setup, placing identical agent and optimizer states on every GPU with unique PRNG keys.
- **Data sharding** for large arrays uses `shmap_vmap` and `shmap_map` from [`evorl/distributed/sharding.py`](https://github.com/emi-group/evorl/blob/main/evorl/distributed/sharding.py) to partition workloads across the device mesh.
- **Gradient averaging** ensures parameter synchronization through `pmean` inside `gradient_update`, maintaining consistency across all devices.

## Frequently Asked Questions

### How does EvoRL detect multiple GPUs automatically?

The training script [`scripts/train.py`](https://github.com/emi-group/evorl/blob/main/scripts/train.py) queries `jax.devices()` at launch. If the list contains more than one device, it sets `enable_multi_devices=True` in the configuration, triggering the `pmap` transformation of the workflow's step and evaluate methods.

### What is the difference between `pmap` and `shmap_vmap` in EvoRL?

`pmap` (used via `RLWorkflow`) replicates the entire computation across devices, suitable for data-parallel RL where each GPU runs a full copy of the environment and agent. `shmap_vmap` (from [`evorl/distributed/sharding.py`](https://github.com/emi-group/evorl/blob/main/evorl/distributed/sharding.py)) partitions large data structures across devices, enabling single-program multiple-data (SPMD) execution for population-based algorithms where the data itself is too large for one device.

### Do I need to modify my algorithm code to support multi-GPU training?

No. As implemented in [`evorl/workflows/rl_workflow.py`](https://github.com/emi-group/evorl/blob/main/evorl/workflows/rl_workflow.py), the workflow API remains identical between single-device and multi-device modes. The framework handles state replication, key splitting, and gradient reduction internally. You only need to toggle `enable_multi_devices` in the Hydra configuration or rely on automatic detection.

### How are random keys managed to ensure reproducibility across GPUs?

The `split_key_to_devices` function in [`evorl/distributed/comm.py`](https://github.com/emi-group/evorl/blob/main/evorl/distributed/comm.py) splits the master PRNG key into distinct sub-keys, one per device. This ensures that while parameters remain synchronized via gradient averaging, the environment randomness and action sampling differ across GPUs, providing diverse data for training.