Lightning-AI / Lightning-AI/pytorch-lightning
NCCL timeout (or GPU OOMs) when using Wandb + configure_model with passing a factory + save_hyperparameters + large models
Nobody has claimed this yet.
- 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
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 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