Lightning-AI / Lightning-AI/pytorch-lightning

NCCL timeout (or GPU OOMs) when using Wandb + configure_model with passing a factory + save_hyperparameters + large models

Open
#19,256 3 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug logger: wandb ver: 2.1.x
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Bug description

Hello! I found this weird interaction that took me a while to debug, so hopefully someone finds it useful or it's possible to fix something in Lightning.

When constructing large models, it's recommended to use `configure_model`. To configure the model creation outside the `lightning`, I've been using factories, so that a fully configured factory can just make a model under the strategy context (e.g. deepspeed).

``` python
def configure_model(self) -> None:
if self.model is None:
self.model = self.model_factory() # make a large model
```

Additionally, I've been using `self.save_hyperparameters()` and `wandb` logger for convenience.
I found that after certain model size, my setup started hanging. I found that the `_sanitize_callable_params` function inside `log_hyperparams` of `WandbLogger` calls my factory again hence temporarily creating yet another copy of a model.

I can't quite find docs on callable parameters for Modules. Is it a bug or a feature? Why would one resolve the callable second time?

Workaround:
`self.save_hyperparameters(ignore="model_factory")`

### What version are you seeing the problem on?

v2.1

### How to reproduce the bug

```python
class MyModule(LightningModule):
def __init__(
self,
model_factory,
**kwargs
):
super().__init__()
self.save_hyperparameters()
self.model = None
self.model_factory = model_factory

def configure_model(self) -> None:
if self.model is None:
self.model = self.model_factory() # make a large model

trainer = Trainer(logger=pl.loggers.WandbLogger(...))
trainer.fit(MyModule(), dataloaders)
```

### Error messages and logs

NCCL hanging for me because rank0 GPU reaches 99% capacity:
```
Unhandled std::runtime_error exception:
[Rank 5] NCCL watchdog thread terminated with exception: [Rank 5] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=3864, OpType=ALLREDUCE, NumelIn=1, NumelOut=1, Timeout(ms)=1800000) ran for 1800838 milliseconds before timing out.
/usr/lib/x86_64-linux-gnu/libstdc++.so.6(+0xd6df4) [0x1513bd4e3df4]
/usr/lib/x86_64-linux-gnu/libpthread.so.0(+0x8609) [0x1514006ac609]
/usr/lib/x86_64-linux-gnu/libc.so.6(clone+0x43) [0x15140046b133]
3bf14e68db564437a3846c1f8ce080e500003N:568:1667 [0] NCCL INFO comm 0x55cbb7249fb0 rank 4 nranks 8 cudaDev 4 busId b00000 - Abort COMPLETE
[E ProcessGroupNCCL.cpp:489] Some NCCL operations have failed or timed out. Due to the asynchronous nature of CUDA kernels, subsequent GPU operations might run on corrupted/incomplete data.
[E ProcessGroupNCCL.cpp:495] To avoid data inconsistency, we are taking the entire process down.
[E ProcessGroupNCCL.cpp:916] [Rank 4] NCCL watchdog thread terminated with exception: [Rank 4] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=3864, OpType=ALLREDUCE, NumelIn=1, NumelOut=1, Timeout(ms)=1800000) ran for 1800826 milliseconds before timing out.
Unhandled std::runtime_error exception:
```
but probably can lead to OOMs?

### Environment

Current environment

```
#- Lightning Component
LightningModule, WandbLogger
#- PyTorch Lightning Version (e.g., 1.5.0):
2.1.2
#- PyTorch Version (e.g., 2.0):
2.1.2
#- Python version (e.g., 3.9):
3.9
```

### More info

_No response_

cc @lantiga @morganmcg1 @borisdayma @scottire @parambharat

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 WandbLogger.log_hyperparams and its _sanitize_callable_params handling, then trace how save_hyperparameters data reaches it during configure_model. Reproduce the example with a factory and a large model under the WandbLogger and distributed setup. Done should include confirmed callable behavior that no longer creates an unintended extra model, plus coverage of the reported interaction.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
distributed-systems, machine-learning, observability
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.