Lightning-AI / Lightning-AI/pytorch-lightning
Support DDP with TPUs
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Description & Motivation
XLA recently introduced support for `DistributedDataParallel` over the native DataParallel approach that we currently support.
The implementation is straightforward. Docs are in https://pytorch.org/xla/release/2.0/index.html#how-to-do-distributeddataparallel and https://github.com/pytorch/xla/blob/master/docs/ddp.md
### Pitch
`Trainer(accelerator="xla", strategy="ddp")` (or `Fabric`) is supported.
### Alternatives
Not do it.
### Additional context
This feature is still experimental and under active development. The performance of the feature doesn't seem to be there.
It requires using PjRT with at least v4.
This issue is meant to plant the idea and track progress if ever made.
cc @borda @carmocca @JackCaoG @steventk-g @Liyang90 @justusschock @awaelchli
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by reading the linked PyTorch/XLA DDP documentation and the Trainer accelerator="xla" strategy="ddp" or Fabric entry points mentioned in the issue. Identify the relevant implementation and tests, then define completion as working DDP support on PjRT v4+ with an appropriate validation path; the issue does not name specific files or tests.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100