Lightning-AI / Lightning-AI/pytorch-lightning
What's the intended way of resuming training on a SLURM cluster?
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### 📚 Documentation
Hi, I'm not sure if this is the intended type of issue for this category, but I thought trying doesn't hurt:
I'm trying to use lightning to train my model on a SLURM cluster due to the high memory requirements. For fairness, it only allows to train 48 hours at once, so I looked up the documentation on how to properly use checkpointing to resume where I left off [as described here](https://pytorch-lightning.readthedocs.io/en/stable/common/checkpointing_basic.html).
The documentation states this:
> ### Resume training state
>
> If you don’t just want to load weights, but instead restore the full training, do the following:
>
> ```
> model = LitModel()
> trainer = Trainer()
>
> # automatically restores model, epoch, step, LR schedulers, etc...
> trainer.fit(model, ckpt_path="some/path/to/my_checkpoint.ckpt")
> ```
So I only let it run for the maximum amount of epochs possible before the wall time would kick in. Lightning then automatically saves a checkpoint in the `lightning_logs/version_21716264/checkpoints/` directory, where `21716264` is the id of the SLURM job that was used to run training. Then when I manually re-queue the job I pass the checkpoint created by the last run via CLI where it is passed to `trainer.fit`. So this seems to work at first glance, but there's something I noticed that seemed odd. When I do this the following warning is issued:
```
Restoring states from the checkpoint path at lightning_logs/version_21716264/checkpoints/epoch=13-step=188874.ckpt
/path/to/venv/lib64/python3.9/site-packages/pytorch_lightning/callbacks/model_checkpoint.py:346: UserWarning: The dirpath has changed from '/path/to/lightning_logs/version_21716264/checkpoints' to '/path/to/lightning_logs/version_21719253/checkpoints', therefore `best_model_score`, `kth_best_model_path`, `kth_value`, `last_model_path` and `best_k_models` won't be reloaded. Only `best_model_path` will be reloaded.
warnings.warn(
```
This seems to suggest that using my approach it "reverts back" to the best model for further training instead of using the last model, which makes training rather ineffective at later stages when validation loss does not always decrease in a couple of epochs.
I tried searching the documentation for a solution to this problem, however it doesn't seem to clearly explain what's the intended way to achieve this. The fact that it uses the SLURM batch id for the version folder seems to suggest, that by default lightning is SLURM aware, so I searched for the lightning documentation on SLURM. I stumbled upon #13773 where I realized a "hpc" keyword exists for the checkpoint, which is supposed to restore the latest "hpc" state. I tried using it, but I just get
```
ValueError: `.fit(ckpt_path="hpc")` is set but no HPC checkpoint was found. Please pass an exact checkpoint path to `.{fn}(ckpt_path=...)`
```
So apparently there is some required prequisite for this? The "hpc" keyword is not mentioned on the [cluster (advanced) page](https://pytorch-lightning.readthedocs.io/en/stable/clouds/cluster_advanced.html), and the [trainer page](https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html) only mentions this as part of the [class API](https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html#trainer-class-api). Nowhere to be seen in samples, but the documentation states it "just works", when it clearly doesn't. Disclaimer at this point, I'm currently using version `1.8.1`, but the changelog doesn't seem to suggest a bugfix in this direction, so I assume it still applies.
So at this point I'm confused. Do I need to store the checkpoint differently? Can't I use the automatic checkpoints for this? Or do I need to override the lightning log directory so it all happens in the same directory so the warning doesn't appear? If so, how? The trainer and SLURM doc pages don't mention anything in this direction. Do I need to define a custom [`ModelCheckpoint`](https://pytorch-lightning.readthedocs.io/en/stable/api/pytorch_lightning.callbacks.ModelCheckpoint.html) callback, or should I define a custom logger, just to use a fixed directory depending on the model? Is there a bug in lightning? I honestly have no idea what I'm supposed to do. The SLURM documentation also mentions auto-requeuing, but I assume that because of my setup bash doesn't pass the signal properly to python. I really hope that's not a precondition for this to work, because _again_ the documentation does not mention anything like this.
Anyways, I feel like there is a straightforward solution to all of my problems the documentation is not telling me. I'm using very limited lightning functionality, so I'm naively assuming it'd just work out of the box:
```python
dm = MyDataModule()
model = MyModel(learning_rate)
trainer = Trainer(
accelerator='gpu' if gpu_count > 0 else None,
devices=gpu_count or None,
max_epochs=max_epochs,
strategy=None if gpu_count <= 1 else DDPStrategy(find_unused_parameters=False, static_graph=True)
)
trainer.fit(model, dm, ckpt_path=checkpoint)
```
I'd be happy if you could point me in the right direction and add your suggestions to the official documentation so that future people won't have the same struggle as I did.
cc @borda @awaelchli
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 by comparing the linked checkpointing, cluster (advanced), Trainer, and ModelCheckpoint documentation with the SLURM example and warning in this issue. Trace how the `ckpt_path` and `hpc` options are described, including checkpoint directories and auto-requeuing. Done means the documentation clearly states the intended resume workflow, prerequisites, and whether automatic checkpoints are sufficient.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- distributed-systems, documentation, machine-learning
- Issue type
- Documentation
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 30/100