Lightning-AI / Lightning-AI/pytorch-lightning

Trainer.save_checkpoint() occasionally saves corrupt checkpoint in Docker/WSL

Open
#21,431 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug checkpointing dockers ver: 2.5.x
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

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 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.