Lightning-AI / Lightning-AI/pytorch-lightning

Share/collect non-tensor objects between processes.

Open
#14,362 5 comments 0 reactions 1 assignee View on GitHub

@dalek-who is already working on this.

Since Apr 17, 2023.

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

Description

## 🚀 Feature

### Motivation

Sometimes we need to share/collect some non-tensor objects between process in DDP mode, mainly includes two kinds of scenes:
- **share status**: for example, a random port to connect database, a timestamp when the experiment starts, an auto-created output directory, etc. These status are created during running and cannot be assigned manually.
- **collect outputs**: for example, in test/predict, collect all outputs to `gloabl_rank_zero` process in `on_{test,preidct}_{batch,epoch}_end` to write predict results and compute metrics over the full dataset.

Current functions cannot satisfy these requests:
- `LightningModule.all_gather` can only deal with tensors. Other objects (like string, id list, Path ...) cannot be gathered or shared.
- `DP` or `DDPSpawn`: they are not as "convenient" as DDP. DP is not encouraged by torch, and DDPSpawn sometimes cannot be started.
- `torch.multiprocessing.SimpleQueue`: there are no tutorials on how to use it with pytorch-lightning DDP mode.

### Pitch

### Alternatives

Consider something like this:
```python
def collect_all_objects(obj: object, global_zero_only: bool=False, default=None):
...
objects_from_process = {
0: obj_0,
1: obj_1,
global_rank_n: obj_n,
}
if not trainer.is_global_zero and global_zero_only:
return default
else:
return objects_from_process
```
This function can also be used to share status between processes: each process collect status and only use the status from rank_0.

**Important:**
When using DDP, the sampler may padding some examples in the last batch and cause the following problem:
- #11102

In this case, manually deduplication is necessary after collecting objects. However, deduplication is much more easier than dealing with inter-process synchronization.

### Additional context

There are some questions related to these requests:

1.**share status**:
- [Share state between DDP processes](https://forums.pytorchlightning.ai/t/share-state-between-ddp-processes/963)
- [Sync output dir between DDP processes](https://forums.pytorchlightning.ai/t/sync-output-dir-between-ddp-processes/779)

2.**collect outputs**:
- #5788
- #10618
- #11086
- #9259

cc @borda @awaelchli @rohitgr7 @akihironitta

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.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.