Lightning-AI / Lightning-AI/pytorch-lightning
AsyncCheckpointIO Should Clone() to CPU not on GPU
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
In [pytorch-lightning/src/lightning/pytorch/plugins/io/async_plugin.py](https://github.com/Lightning-AI/pytorch-lightning/blame/master/src/lightning/pytorch/plugins/io/async_plugin.py#L100)
Tensor in Model Checkpoint are cloned() to prevent race condition when doing async uploading in threads.
```
# snapshot the checkpoint payload on the caller thread to avoid races with parameter mutation
def _clone_tensor(t: torch.Tensor) -> torch.Tensor:
"""Clones a tensor on the caller thread."""
# detach to avoid autograd history and clone to take a point-in-time copy
return t.detach().clone()
```
However, here it is clone() from GPU memory to GPU memory. Given GPU memory is often limited, this step is dangerous when num_thread go up.
A more clean solution is to clone the tensors to CPU. This achieves the same purpose without using without using GPU memory. CPU memory is abundant most of the time.
```
# snapshot the checkpoint payload on the caller thread to avoid races with parameter mutation
def _clone_tensor(t: torch.Tensor) -> torch.Tensor:
"""Clones a tensor on the caller thread."""
# detach to avoid autograd history and clone to take a point-in-time copy
return t.detach().cpu().clone()
```
### What version are you seeing the problem on?
master
### Reproduced in studio
_No response_
### How to reproduce the bug
```python
```
### Error messages and logs
If cloning on GPU memory, this is dangerous for large model checkpoint (e.g. 15GB in our case)
Cloning to CPU
```
[ASYNC CHECKPOINT BEFORE clone] GPU 0: allocated=21.54 GB, reserved=124.30 GB
[ASYNC CHECKPOINT AFTER clone] GPU 0: allocated=21.54 GB, reserved=124.30 GB
```
Cloning on GPU
```
[ASYNC CHECKPOINT BEFORE clone] GPU 0: allocated=21.54 GB, reserved=124.30 GB
[ASYNC CHECKPOINT AFTER clone] GPU 0: allocated=37.54 GB, reserved=124.30 GB
```
### Environment
Current environment
```
#- PyTorch Lightning Version (e.g., 2.6.0):
#- PyTorch Version (e.g., 2.5):
#- 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 @ethanwharris
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 in src/lightning/pytorch/plugins/io/async_plugin.py at _clone_tensor, where checkpoint tensors are detached and cloned. Verify the snapshot moves tensor data to CPU without retaining an extra GPU copy, while preserving the race-avoidance behavior needed by asynchronous uploads. Done means the async checkpoint path no longer duplicates checkpoint tensors in GPU memory.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning, performance
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 55/100