Lightning-AI / Lightning-AI/pytorch-lightning

ModelCheckpoint save ckpts at the end of every epoch even in step-saving strategy

Open
#20,075 0 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

bug callback: model checkpoint ver: 2.2.x
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Bug description

When `ModelCheckpoint` is enabled and **no validation set is given**, the trainer still save the model checkpoint at the end of every epoch.

I have turn on the `every_n_train_steps=1000` and `every_n_epochs=None`. And the expected behavior according to the docs of [ModelCheckpoint API](https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.ModelCheckpoint.html#lightning.pytorch.callbacks.ModelCheckpoint) should **only save the checkpoints at every 1000 steps, and should not save at the end of each epoch**.

This slows down the training process especially when the training dataset is small. And it has cost me 10TBW of my SSDs.

A quick solution is to pass `save_on_train_epoch_end=False` to the `ModelCheckpoint`.

I believe this bug is caused by the failure of checking the `_should_save_on_train_epoch_end` at L320.
https://github.com/Lightning-AI/pytorch-lightning/blob/bf25167bbf64f50ba335aa759318946b21775cd2/src/lightning/pytorch/callbacks/model_checkpoint.py#L318-L324

And the logic at L432-437 triggers the bug when **no validation dataset is given**.
https://github.com/Lightning-AI/pytorch-lightning/blob/bf25167bbf64f50ba335aa759318946b21775cd2/src/lightning/pytorch/callbacks/model_checkpoint.py#L423-L441

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

master

### How to reproduce the bug

```python
Run lightning with following callback, with a relative small dataset (e.g., a dataset with only 100 batch in one epoch)

ModelCheckpoint(
monitor="train_loss",
every_n_train_steps=10,
dirpath="checkpoints",
filename="{step}-{train_loss:.3f}",
save_top_k=1,
mode="min",
save_last=True,
)
```

### Error messages and logs

_No response_

### Environment

Current environment

* CUDA:
- GPU:
- NVIDIA A800 80GB PCIe
- NVIDIA A800 80GB PCIe
- NVIDIA A800 80GB PCIe
- NVIDIA A800 80GB PCIe
- available: True
- version: 12.1
* Lightning:
- lightning: 2.3.0
- lightning-utilities: 0.11.2
- pytorch-lightning: 2.3.0
- torch: 2.3.1+cu121
- torchaudio: 2.3.1+cu121
- torchmetrics: 1.4.0.post0
- torchvision: 0.18.1+cu121
* Packages:
- absl-py: 2.1.0
- aiohttp: 3.9.5
- aiosignal: 1.3.1
- annotated-types: 0.7.0
- antlr4-python3-runtime: 4.9.3
- anyio: 4.4.0
- argon2-cffi: 23.1.0
- argon2-cffi-bindings: 21.2.0
- arrow: 1.3.0
- asttokens: 2.4.1
- async-lru: 2.0.4
- attrs: 23.2.0
- babel: 2.15.0
- beautifulsoup4: 4.12.3
- bleach: 6.1.0
- certifi: 2024.6.2
- cffi: 1.16.0
- cfgv: 3.4.0
- charset-normalizer: 3.3.2
- comm: 0.2.2
- contourpy: 1.2.1
- cycler: 0.12.1
- debugpy: 1.8.1
- decorator: 5.1.1
- deepspeed: 0.14.4
- defusedxml: 0.7.1
- dill: 0.3.8
- distlib: 0.3.8
- einops: 0.8.0
- executing: 2.0.1
- fastjsonschema: 2.20.0
- filelock: 3.15.1
- flash-attn: 2.5.9.post1
- fonttools: 4.53.0
- fqdn: 1.5.1
- frozenlist: 1.4.1
- fsspec: 2024.6.0
- grpcio: 1.64.1
- h11: 0.14.0
- hjson: 3.1.0
- httpcore: 1.0.5
- httpx: 0.27.0
- huggingface-hub: 0.23.4
- hydra-core: 1.3.2
- identify: 2.5.36
- idna: 3.7
- ipykernel: 6.29.4
- ipython: 8.25.0
- isoduration: 20.11.0
- jedi: 0.19.1
- jinja2: 3.1.4
- json5: 0.9.25
- jsonpointer: 3.0.0
- jsonschema: 4.22.0
- jsonschema-specifications: 2023.12.1
- jupyter-client: 8.6.2
- jupyter-core: 5.7.2
- jupyter-events: 0.10.0
- jupyter-lsp: 2.2.5
- jupyter-server: 2.14.1
- jupyter-server-terminals: 0.5.3
- jupyterlab: 4.2.2
- jupyterlab-pygments: 0.3.0
- jupyterlab-server: 2.27.2
- kiwisolver: 1.4.5
- lightning: 2.3.0
- lightning-utilities: 0.11.2
- markdown: 3.6
- markupsafe: 2.1.5
- matplotlib: 3.9.0
- matplotlib-inline: 0.1.7
- mistune: 3.0.2
- mpmath: 1.3.0
- multidict: 6.0.5
- nbclient: 0.10.0
- nbconvert: 7.16.4
- nbformat: 5.10.4
- nest-asyncio: 1.6.0
- networkx: 3.3
- ninja: 1.11.1.1
- nodeenv: 1.9.1
- notebook-shim: 0.2.4
- numpy: 1.26.4
- nvidia-cublas-cu12: 12.1.3.1
- nvidia-cuda-cupti-cu12: 12.1.105
- nvidia-cuda-nvrtc-cu12: 12.1.105
- nvidia-cuda-runtime-cu12: 12.1.105
- nvidia-cudnn-cu12: 8.9.2.26
- nvidia-cufft-cu12: 11.0.2.54
- nvidia-curand-cu12: 10.3.2.106
- nvidia-cusolver-cu12: 11.4.5.107
- nvidia-cusparse-cu12: 12.1.0.106
- nvidia-ml-py: 12.555.43
- nvidia-nccl-cu12: 2.20.5
- nvidia-nvjitlink-cu12: 12.5.40
- nvidia-nvtx-cu12: 12.1.105
- omegaconf: 2.3.0
- overrides: 7.7.0
- packaging: 24.1
- pandas: 2.2.2
- pandocfilters: 1.5.1
- parso: 0.8.4
- pexpect: 4.9.0
- pillow: 10.3.0
- pip: 24.0
- platformdirs: 4.2.2
- pre-commit: 3.7.1
- prometheus-client: 0.20.0
- prompt-toolkit: 3.0.47
- protobuf: 4.25.3
- psutil: 6.0.0
- ptyprocess: 0.7.0
- pure-eval: 0.2.2
- py-cpuinfo: 9.0.0
- pycparser: 2.22
- pydantic: 2.7.4
- pydantic-core: 2.18.4
- pyfunctional: 1.5.0
- pygments: 2.18.0
- pyparsing: 3.1.2
- python-dateutil: 2.9.0.post0
- python-json-logger: 2.0.7
- pytorch-lightning: 2.3.0
- pytz: 2024.1
- pyyaml: 6.0.1
- pyzmq: 26.0.3
- referencing: 0.35.1
- regex: 2024.5.15
- requests: 2.32.3
- rfc3339-validator: 0.1.4
- rfc3986-validator: 0.1.1
- rpds-py: 0.18.1
- safetensors: 0.4.3
- send2trash: 1.8.3
- setuptools: 69.5.1
- six: 1.16.0
- sniffio: 1.3.1
- soupsieve: 2.5
- stack-data: 0.6.3
- sympy: 1.12.1
- tabulate: 0.9.0
- tensorboard: 2.17.0
- tensorboard-data-server: 0.7.2
- tensorboardx: 2.6.2.2
- termcolor: 2.4.0
- terminado: 0.18.1
- tinycss2: 1.3.0
- tokenizers: 0.19.1
- torch: 2.3.1+cu121
- torchaudio: 2.3.1+cu121
- torchmetrics: 1.4.0.post0
- torchvision: 0.18.1+cu121
- tornado: 6.4.1
- tqdm: 4.66.4
- traitlets: 5.14.3
- transformers: 4.42.3
- triton: 2.3.1
- types-python-dateutil: 2.9.0.20240316
- typing-extensions: 4.12.2
- tzdata: 2024.1
- uri-template: 1.3.0
- urllib3: 2.2.2
- virtualenv: 20.26.2
- wcwidth: 0.2.13
- webcolors: 24.6.0
- webencodings: 0.5.1
- websocket-client: 1.8.0
- werkzeug: 3.0.3
- wheel: 0.43.0
- yarl: 1.9.4
* System:
- OS: Linux
- architecture:
- 64bit
- ELF
- processor: x86_64
- python: 3.11.9
- release: 5.15.0-67-generic
- version: #74~20.04.1-Ubuntu SMP Wed Feb 22 14:52:34 UTC 2023

### More info

_No response_

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 in src/lightning/pytorch/callbacks/model_checkpoint.py at the referenced lines around _should_save_on_train_epoch_end and the logic for runs without a validation dataset. Reproduce the issue with every_n_train_steps=10, every_n_epochs=None, and a small dataset. Done means checkpoints are saved at the configured step intervals without an additional checkpoint at the end of every epoch.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Clearly specified
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.