Lightning-AI / Lightning-AI/pytorch-lightning
Trainer.save_checkpoint() occasionally saves corrupt checkpoint in Docker/WSL
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
An exception may occur when loading a model checkpoint (with `LightningModule.load_from_checkpoint()`) that was corrupted during saving with PyTorch Lightning (`Trainer.save_checkpoint()`).
Trying to load the checkpoint file directly e.g. with `torch.load()` shows the same exception. Comparing the file with a valid checkpoint (e.g. with `diff`) confirms that both differ. This indicates that the checkpoint file is corrupted during saving.
This exception appears for CPU and CUDA versions of PyTorch libraries. Tested with PyTorch 2.9.0+cu130 / Torchvision 0.24.0+cu130 (for CUDA 13.0) and the CPU equivalent, and lightning-2.6.0 (bug initially spotted with lightning-2.5.5).
A minimal self-contained script to reproduce the issue is attached below. Continuously saving and loading a checkpoint after `trainer.fit()` will sometimes raise an exception and other times succeed. Script output might look like this:
```
Starting loop...
10 ok, 0 fail
20 ok, 0 fail
30 ok, 0 fail
40 ok, 0 fail
48 ok, 1 fail: PytorchStreamReader failed reading file data/1017: invalid header or archive is corrupted
48 ok, 2 fail: PytorchStreamReader failed reading file data/569: invalid header or archive is corrupted
50 ok, 2 fail
60 ok, 2 fail
```
Failure also occurs with `trainer.save_checkpoint(weights_only=True)`.
On an older version of PyTorch (torch-2.3.1+cu121, torchvision-0.18.1+cu121) the issue is even more prevalent with an approx. 1/6 failure rate.
On the older PyTorch version, sleeping between saving and loading increases the likelihood of failure; with `sleep(10)` every save is corrupted. On the newer PyTorch version, no such effect can be observed.
Test script is running in a docker container on Windows/WSL. Failures occur regardless of storing the checkpoint on a bind mount or volume mount.
Issue cannot be reproduced on another machine with the same docker image, running on Linux.
Issue cannot be reproduced with Lightning Studio: https://lightning.ai/stf9790/vision-model/studios/prime-bronze-7x9h/code?source=copylink
### What version are you seeing the problem on?
v2.6
### Reproduced in studio
Not reproducible in studio.
### How to reproduce the bug
```python
from time import sleep
import traceback
import torch
import torch.nn as nn
import torchvision
from torchvision.transforms import v2
import lightning as L
class MyImageClassifier(L.LightningModule):
def __init__(self, num_classes: int):
super().__init__()
self.loss_function = nn.CrossEntropyLoss()
self.model = self.instantiate_model(num_classes)
self.save_hyperparameters()
def instantiate_model(self, num_classes: int) -> nn.Module:
model = torchvision.models.maxvit_t(weights='DEFAULT')
block_channels=[64, 128, 256, 512]
model.classifier = nn.Sequential(
nn.AdaptiveAvgPool2d(1),
nn.Flatten(),
nn.LayerNorm(block_channels[-1]),
nn.Linear(block_channels[-1], block_channels[-1]),
nn.Tanh(),
nn.Linear(block_channels[-1], num_classes, bias=False)
)
return model
def configure_optimizers(self):
optimizer = torch.optim.SGD(self.model.parameters(), lr = 0.001, momentum=0.9)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer = optimizer, step_size = 2, gamma = 0.75)
return {
'optimizer': optimizer,
'lr_scheduler': scheduler,
}
def forward(self, inputs):
return self.model(inputs)
def training_step(self, batch, batch_idx, dataloader_idx=0):
inputs, targets = batch
logits = self(inputs)
loss = self.loss_function(logits, targets)
return loss
def validation_step(self, batch, batch_idx, dataloader_idx=0):
inputs, targets = batch
logits = self(inputs)
loss = self.loss_function(logits, targets)
return loss
def worker():
transform = v2.Compose([
v2.ToImage(),
v2.ToDtype(torch.float32, scale=True),
])
dataset = torchvision.datasets.fakedata.FakeData(size=100, num_classes=3, transform=transform)
train_dataloader = torch.utils.data.DataLoader(dataset, batch_size=16)
val_dataloader = torch.utils.data.DataLoader(dataset, batch_size=16)
model = MyImageClassifier(num_classes=3)
trainer = L.Trainer(
max_epochs=1,
# accelerator='cpu',
)
trainer.fit(
model=model,
train_dataloaders=train_dataloader,
val_dataloaders=val_dataloader,
)
ok = 0
fail = 0
print("Starting loop...")
while True:
trainer.save_checkpoint('tmp_test.ckpt', weights_only=False)
# sleep(10)
try:
newnet = MyImageClassifier.load_from_checkpoint(checkpoint_path='tmp_test.ckpt')
ok += 1
except Exception as e:
fail += 1
print(f"{ok} ok, {fail} fail: {str(e)}")
#print(traceback.format_exc())
if ok % 10 == 0:
print(f"{ok} ok, {fail} fail")
if __name__ == "__main__":
worker()
```
### Error messages and logs
Output might look similar to this. Exceptions occur randomly, and the `data/` may vary for each exception.
```
GPU available: False, used: False
TPU available: False, using: 0 TPU cores
/root/.cache/pypoetry/virtualenvs/test-s5OTQf_9-py3.12/lib/python3.12/site-packages/lightning/pytorch/trainer/connectors/logger_connector/logger_connector.py:76: Starting from v1.9.0, `tensorboardX` has been removed as a dependency of the `lightning.pytorch` package, due to potential conflicts with other packages in the ML ecosystem. For this reason, `logger=True` will use `CSVLogger` as the default logger, unless the `tensorboard` or `tensorboardX` packages are found. Please `pip install lightning[extra]` or one of them to enable TensorBoard support by default
| Name | Type | Params | Mode | FLOPs
-------------------------------------------------------------------
0 | loss_function | CrossEntropyLoss | 0 | train | 0
1 | model | MaxVit | 30.4 M | train | 0
-------------------------------------------------------------------
30.4 M Trainable params
0 Non-trainable params
30.4 M Total params
121.637 Total estimated model params size (MB)
671 Modules in train mode
0 Modules in eval mode
0 Total Flops
Sanity Checking: | | 0/? [00:00
Current environment
* CUDA:
- GPU: None
- available: False
- version: None
* Lightning:
- lightning: 2.6.0
- lightning-utilities: 0.15.2
- pytorch-lightning: 2.6.0
- torch: 2.9.0+cpu
- torchmetrics: 1.8.2
- torchvision: 0.24.0+cpu
* Packages:
- aiohappyeyeballs: 2.6.1
- aiohttp: 3.13.2
- aiosignal: 1.4.0
- attrs: 25.4.0
- autocommand: 2.2.2
- backports.tarfile: 1.2.0
- filelock: 3.20.1
- frozenlist: 1.8.0
- fsspec: 2025.12.0
- idna: 3.11
- importlib-metadata: 8.0.0
- inflect: 7.3.1
- jaraco.collections: 5.1.0
- jaraco.context: 5.3.0
- jaraco.functools: 4.0.1
- jaraco.text: 3.12.1
- jinja2: 3.1.6
- lightning: 2.6.0
- lightning-utilities: 0.15.2
- markupsafe: 3.0.3
- more-itertools: 10.3.0
- mpmath: 1.3.0
- multidict: 6.7.0
- networkx: 3.6.1
- numpy: 2.3.5
- packaging: 25.0
- pillow: 12.0.0
- pip: 25.3
- platformdirs: 4.2.2
- propcache: 0.4.1
- pytorch-lightning: 2.6.0
- pyyaml: 6.0.3
- setuptools: 80.9.0
- sympy: 1.14.0
- tomli: 2.0.1
- torch: 2.9.0+cpu
- torchmetrics: 1.8.2
- torchvision: 0.24.0+cpu
- tqdm: 4.67.1
- typeguard: 4.3.0
- typing-extensions: 4.15.0
- wheel: 0.45.1
- yarl: 1.22.0
- zipp: 3.19.2
* System:
- OS: Linux
- architecture:
- 64bit
-
- processor: x86_64
- python: 3.12.3
- release: 6.6.87.2-microsoft-standard-WSL2
- version: #1 SMP PREEMPT_DYNAMIC Thu Jun 5 18:30:46 UTC 2025
### More info
_No response_
cc @ethanwharris @lantiga
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 Trainer.save_checkpoint(), then trace the checkpoint through torch.save/torch.load and LightningModule.load_from_checkpoint() using the provided Docker/WSL reproducer. Compare behavior across the reported PyTorch versions and storage environments. Done means the repeated save-and-load test no longer produces corrupt checkpoints in the affected environment.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100