Support custom process group backends

[tchaton]

🚀 Feature

Motivation

Find the link there: https://github.com/facebookresearch/fairring

Pitch

Alternatives

Additional context

---

If you enjoy Lightning, check out our other projects! ⚡

• Metrics: Machine learning metrics for distributed, scalable PyTorch applications.

• Lite: enables pure PyTorch users to scale their existing code on any kind of device while retaining full control over their own loops and optimization logic.

• Flash: The fastest way to get a Lightning baseline! A collection of tasks for fast prototyping, baselining, fine-tuning, and solving problems with deep learning.

• Bolts: Pretrained SOTA Deep Learning models, callbacks, and more for research and production with PyTorch Lightning and PyTorch.

• Lightning Transformers: Flexible interface for high-performance research using SOTA Transformers leveraging Pytorch Lightning, Transformers, and Hydra.

cc @Borda @awaelchli @rohitgr7 @akihironitta

[ananthsub]

I've thought about this as well. IMO, the process group backend should be an optional constructor argument on the relevant strategies. This would make it super simple to integrate with libraries like fairring, but still allow Lightning to set a good default depending on the device type being used.

e.g.

1. On the relevant strategies, add process_group_backend (or pg_backend as a shorter name) to the constructor:

class DDPStrategy(ParallelStrategy):
    def __init__(..., pg_backend: Optional[str] = None):
         self._pg_backend: Optional[str] = pg_backend

2. Add a utility function to distributed.py that gets the default process group backend for a given device. this is purely to facilitate code reuse across strategies instead of needing to inherit or extend from ParallelStrategy.

# Utility function that can be shared across strategies
def get_default_process_group_backend_for_device(device: torch.device) -> str:
    return "nccl" if device.type == "cuda" else "gloo"

3. Inside of setup_distributed make this change from

init_dist_connection(self.cluster_environment, self.torch_distributed_backend)

to

self._pg_backend = self._pg_backend or get_default_process_group_backend_for_device(self.root_device)
init_dist_connection(self.cluster_environment, self._pg_backend)

4. Deprecate/remove this property from the ParallelStrategy along with setting the process group backend via an environment variable. This is super hidden and not documented anywhere. https://github.com/PyTorchLightning/pytorch-lightning/blob/4d72110b51525296eac16454e7b34e2c1b60cb7e/pytorch_lightning/strategies/parallel.py#L101-L106

Then the end user simply does this:

import fairring
Trainer(strategy=DDPStrategy(pg_backend="fairring"), accelerator="gpu", devices=8)

[daniellepintz]

I can work on this! (Making process group backend a constructor argument and then integrating with fairring)

[awaelchli]

The previous mechanism with the env var had the advantage of making the code hardware agnostic, i.e., you could set PL_TORCH_DISTRIBUTED_BACKEND="gloo" as fixed and freely choose the devices and hardware though e.g. CLI. Now, you would need to change the code and instantiate the strategy yourself, or register the custom strategy in the Registry.
I hope existing users will not be impacted by this inconvenience.

An alternative would be to keep the env variable.

[ananthsub]

@awaelchli - if there is pushback from users regarding the enviroment variable, we can always undeprecate it and continue supporting it. At the very least, this logic is now isolated to a single utility function under distributed.py