Lightning-AI / Lightning-AI/pytorch-lightning

Provide a backend agnostic Join for LightningLite

Open
#14,635 4 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

fabric feature
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

## 🚀 Feature

Provide Join through an intuitive API in LightningLite and make it backend agnostic, i.e., switching from DDP to single-device and vice versa should not require changes to the code.

### Motivation

The [DDP Join context manager](https://pytorch.org/tutorials/advanced/generic_join.html) in PyTorch allows you to run your loops with different number of items on each rank, without getting out of sync issues and hangs in collective calls. PyTorch calls this "uneven inputs". Normally, the DistributedSampler would "even out" the data on each rank by inserting fake, repeated data.

### Pitch

Provide Join in LightningLite. More specifically, through the Strategy. The idea here is that once the user added the join to their loop, they won't have to change it again when switching to single-device strategy (it would simply be a no-op).

Note, by default Lite will auto-insert a DistributedSampler on the dataloader for the user. The tricky part here is that Join is only useful if you set `drop_last=False` in the sampler. How do we link the two features together so that they work in a meaningful way?

### Alternatives

Do not introduce this. The user can just use the raw PyTorch APIs.

### Additional context

Once this lands in Lite, the PL strategies can also make use of it in their implementations. This can be developed in parallel to #3325.

______________________________________________________________________

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

- [**Metrics**](https://github.com/Lightning-AI/metrics): Machine learning metrics for distributed, scalable PyTorch applications.

- [**Lite**](https://pytorch-lightning.readthedocs.io/en/latest/starter/lightning_lite.html): 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**](https://github.com/Lightning-AI/lightning-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**](https://github.com/Lightning-AI/lightning-bolts): Pretrained SOTA Deep Learning models, callbacks, and more for research and production with PyTorch Lightning and PyTorch.

- [**Lightning Transformers**](https://github.com/Lightning-AI/lightning-transformers): Flexible interface for high-performance research using SOTA Transformers leveraging PyTorch Lightning, Transformers, and Hydra.

cc @borda @carmocca @justusschock @awaelchli

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with LightningLite's Strategy abstraction and the PyTorch DDP Join context manager described in the issue, then inspect how Lite inserts DistributedSampler instances. Define how drop_last=False interacts with Join and how the API behaves for DDP versus single-device strategies. Done means users can keep the same loop while Join operates for distributed execution and is a no-op on one device.

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
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.