facebookresearch / facebookresearch/fairscale
[feat] Pipe RPC + DDP for pipe parallelism + data parallelism
- Dominant language
- Python
- Stars
- 3.4k
- Forks
- 293
- PR merge metrics
- No merged PRs in 30d
Description
## 🚀 Feature
Use Pipe RPC with DDP! There is also a chance this is supported, and I'm doing things wrong :)
## Motivation
We currently have FairScale pipe parallelism RPC within Lightning. It's easier for the user to use RPC as this allows us to add model parallelism without changing the training loop flow, however I don't think there is support for DDP, as communication is restricted to 'trainer processes' which are responsible for gradient flow to the workers. There seems to be a solution [here](https://pytorch.org/tutorials/advanced/rpc_ddp_tutorial.html) that can act as backbone for this!
## Pitch
Snippet taking the Pipe RPC script:
```python
if rank == 1:
# For RPC, all ranks other than 0 just need to call rpc.shutdown()
torch.distributed.rpc.shutdown()
return
model = nn.Sequential(torch.nn.Linear(10, 10), torch.nn.ReLU(), torch.nn.Linear(10, 5))
target = torch.randint(0, 2, size=(20, 1)).squeeze()
data = torch.randn(20, 10)
loss_fn = F.nll_loss
device = torch.device("cuda", rank)
# Split my model across two GPUs
model = fairscale.nn.PipeRPCWrapper(
model,
balance=[2, 1],
worker_map={0: "worker0", 1: "worker1"}, # Needed to convert ranks to RPC worker names
input_device=device,
).to(device)
# Replicate across my other GPUs
model = torch.nn.DistributedParallel(model, process_group=mpu.get_data_parallel_group())
```
## Alternatives
Force to use the multiprocess version, which is simpler as it doesn't require RPC, but less elegant as a long term solution imo.
## Additional context
If I get some cycles I will also spend time investigating this to see if we can support this. Currently pipeline parallelism with RPC has some really nice features like checkpoint/microbatching, but the lack of DDP makes this much less useable. All the thoughts please!
Contributor guide
Assessment
This issue has not been assessed yet.