Multi-GPU PPO
PPO scaled across multiple GPUs with PyTorch DistributedDataParallel and torchrun.
Pythonopen source
Problem
A single-GPU PPO implementation hits a wall on wall-clock time once environments and networks grow. The aim was to scale the same algorithm horizontally without changing its learning behaviour.
Approach
- Wrapped the policy in PyTorch DistributedDataParallel so gradients are all-reduced across processes automatically on each backward pass.
- Launched with torchrun, one process per GPU, each collecting its own rollouts so experience collection scales with the number of devices.
- Kept the optimisation maths identical to single-GPU PPO, so results stay comparable while throughput increases.
- Handled seeding and logging per rank so runs remain reproducible and only rank zero writes metrics.
Tech
PyTorchDistributedDataParalleltorchrunCUDAPython
Status
Shipped and open source. Built as the distributed counterpart to the from-scratch PPO project.