Lightning-AI / Lightning-AI/pytorch-lightning

OnExceptionCheckpoint callback suppresses exceptions and results in NCCL timeout

Open
#20,187 0 comments 0 reactions 0 assignees View on GitHub
bug checkpointing strategy: ddp ver: 2.4.x
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Bug description

When running Lightning with a multi-device training strategy (e.g. with DDP), using the `OnExceptionCheckpoint` callback:
- silently swallows exceptions, which makes it challenging to identify the cause of errors
- results in a NCCL timeout

This is due to the following:
- When we catch an exception, it gets handled by `_call_and_handle_interrupt`, which calls into `_interrupt`: https://github.com/Lightning-AI/pytorch-lightning/blob/1551a16b94f5234a4a78801098f64d0732ef5cb5/src/lightning/pytorch/trainer/call.py#L67
- We are supposed to re-raise the original exception at the end of this function, but we never get there because...
- In `_interrupt`, we call `_call_callback_hooks`, which calls the `on_exception` callbacks: https://github.com/Lightning-AI/pytorch-lightning/blob/1551a16b94f5234a4a78801098f64d0732ef5cb5/src/lightning/pytorch/trainer/call.py#L76
- If the `OnExceptionCheckpoint` is enabled, we then call that callback. However, we never finish executing this callback, because in that callback, we call `trainer.save_checkpoint`: https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/pytorch/callbacks/on_exception_checkpoint.py#L67
- The `trainer.save_checkpoint` method saves the checkpoint, and then calls `self.strategy.barrier("Trainer.save_checkpoint")`, which waits for the other processes to get reach that barrier. However, if those processes haven't had an exception, they will never hit this codepath, which means we never advance beyond that barrier (until it times out).

As described in the docstring for `Trainer.save_checkpoint`:
> This method needs to be called on all processes in case the selected strategy is handling distributed checkpointing.

In practice, this means that our jobs eventually time out with a NCCL error, and don't print the original exception.

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

v2.4

### How to reproduce the bug

_No response_

### Error messages and logs

```
# Error messages and logs here please
```

### Environment

Current environment

```
#- PyTorch Lightning Version (e.g., 2.4.0):
#- PyTorch Version (e.g., 2.4):
#- Python version (e.g., 3.12):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
```

### More info

_No response_

cc @lantiga @justusschock

Contributor guide

Open the contributing guide

Research direction

Read src/lightning/pytorch/trainer/call.py around _call_and_handle_interrupt and _interrupt, then inspect src/lightning/pytorch/callbacks/on_exception_checkpoint.py and Trainer.save_checkpoint. Reproduce with a multi-device DDP run using OnExceptionCheckpoint. Done means the original exception remains visible and the callback does not cause a NCCL barrier timeout.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
distributed-systems
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.