Lightning-AI / Lightning-AI/pytorch-lightning

Can't resume automatically a job, ckpt_path="hpc" throws ValueError from the start

Open
#20,347 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug checkpointing environment: slurm ver: 2.4.x
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

## Summary

When attempting to resume a job from where it left off before reaching wall-time on a SLURM cluster using PyTorch Lightning, the ckpt_path="hpc" option causes an error if no HPC checkpoint exists yet. This prevents the initial training run from starting.

## Expected Behavior

- The job should be able to resume from an HPC checkpoint if one exists when using in combination:
- `SLURMEnvironment(auto_requeue=True, requeue_signal=signal.SIGUSR1)`
- `trainer.fit(model, datamodule=dm, ckpt_path="hpc")`
- `#SBATCH --signal=SIGUSR1@30`

- If no HPC checkpoint exists (e.g., on the first run), the job should start training from scratch without throwing an error. Currently it throws one:

```
.fit(ckpt_path="hpc")` is set but no HPC checkpoint was found.' Please pass an exact checkpoint path to `.fit(ckpt_path=...)
```

## Current Behavior

### Using `ckpt_path=None` allows the job to start but doesn't resume from the HPC checkpoint when one is created.

If I use `trainer.fit(model, datamodule=dm, ckpt_path=None)`, the SIGUSR1 is correctly catched and the checkpoint `hpc_ckpt_1.ckpt` correctly created. However the checkpoint is not used which is expected because we left `ckpt_path=None`.

```
requeing job 245038...
Requeued SLURM job: 245038
srun: Job step aborted: Waiting up to 62 seconds for job step to finish.
slurmstepd: error: *** JOB 245038 ON jzxh061 CANCELLED AT 2024-10-17T15:45:56 DUE TO JOB REQUEUE ***
slurmstepd: error: *** STEP 245038.0 ON jzxh061 CANCELLED AT 2024-10-17T15:45:56 DUE TO JOB REQUEUE ***
```

### Using `ckpt_path="hpc"` throws an error if no HPC checkpoint is found, preventing the initial training run.

The logic of looking for and loading the hpc checkpoint from what I understood should be handled by setting `ckpt_path="hpc"`
However, as can be seen in https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/pytorch/trainer/connectors/checkpoint_connector.py#L193C1-L199C46
if an hpc ckpt is not found it throws an error and stops:
```
.fit(ckpt_path="hpc")` is set but no HPC checkpoint was found.' Please pass an exact checkpoint path to `.fit(ckpt_path=...)
```
The issue is that for the very first training of course there would be no hpc checkpoint because we haven't started any training yet

### Relevant issues

#16639

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

v2.4

### How to reproduce the bug

dummy_model.py
```python
import os
import torch
import lightning as L
from torch.utils.data import Dataset
from lightning.pytorch.callbacks import ModelCheckpoint
import argparse
from torchdata.stateful_dataloader import StatefulDataLoader
from torch.distributed.checkpoint.stateful import Stateful
from torch.utils.data.distributed import DistributedSampler
import pickle
import signal
from lightning.pytorch.plugins.environments import SLURMEnvironment
import time

class DummyDataset(Dataset):
def __init__(self, size=100000):
self.size = size
self.data = torch.randn(size, 10)
self.labels = torch.randint(0, 2, (size,))
self.current_index = 0

def __len__(self):
return self.size

def __getitem__(self, idx):
print(f"Accessing index: {idx}, {self.data[idx]}")
return self.data[idx], self.labels[idx]

class DPAwareDataLoader(StatefulDataLoader, Stateful):
def __init__(self, dataset: Dataset, batch_size: int, sampler=None, **kwargs):
super().__init__(dataset, batch_size=batch_size, sampler=sampler, **kwargs)
self._rank_id = f"dp_rank_{sampler.rank if sampler else 0}"
print(self._rank_id, " initialized")

def state_dict(self):
print(f"self._rank_id: ", f"{super().state_dict()}")
return {self._rank_id: super().state_dict()}

def load_state_dict(self, state_dict):
if not state_dict:
return
if self._rank_id not in state_dict:
print(f"DataLoader state is empty for dp rank {self._dp_rank}, expected key {self._rank_id}")
return
print(f"DataLoader state loading for dp rank {self._dp_rank}")
super().load_state_dict(state_dict[self._rank_id])

class DummyDataModule(L.LightningDataModule, Stateful):
def __init__(self, batch_size=32):
super().__init__()
self.batch_size = batch_size
self.train_dataset = None
self.dataloader = None

def setup(self, stage=None):
self.train_dataset = DummyDataset()

# DistributedSampler automatically retrieves world_size and rank
# from the current distributed group.
#
# ref: https://pytorch.org/docs/stable/data.html#torch.utils.data.distributed.DistributedSampler
#
# In PyTorch Lightning:
# - The distributed environment is initialized by the Trainer.
# - This sets up the process group with the correct world_size and rank.
# - DistributedSampler then uses these values automatically.
#
# By not specifying num_replicas and rank, we allow DistributedSampler
# to adapt to the current distributed setup, making our code more flexible.
# This works seamlessly with PyTorch Lightning's managed distributed training.
#
# Note: This automatic retrieval only works correctly if the distributed
# environment has been initialized, which Lightning ensures before calling setup().
self.sampler = DistributedSampler(
self.train_dataset,
shuffle=False # Ensure deterministic data order across processes for testing purposes
)

def train_dataloader(self):
if self.dataloader is None:
self.dataloader = DPAwareDataLoader(
self.train_dataset,
batch_size=self.batch_size,
sampler=self.sampler,
num_workers=2
)
return self.dataloader

def state_dict(self):
return {
"dataloader_state": self.dataloader.state_dict() if self.dataloader else None,
}

def load_state_dict(self, state_dict):
if self.dataloader and state_dict["dataloader_state"]:
self.dataloader.load_state_dict(state_dict["dataloader_state"])

class DummyModel(L.LightningModule):
def __init__(self):
super().__init__()
self.layer = torch.nn.Linear(10, 2)
self.loss = torch.nn.CrossEntropyLoss()
self.example_count = 0
self.custom_global_step = 0

def training_step(self, batch, batch_idx):
time.sleep(10)
x, y = batch
y_hat = self.layer(x)
loss = self.loss(y_hat, y)
self.log('train_loss', loss)

self.example_count += len(x)
self.custom_global_step = self.global_step

if self.example_count % 100 == 0:
print(f"GPU {self.global_rank}: Processed {self.example_count} examples, Global Step: {self.custom_global_step}")

return loss

def on_save_checkpoint(self, checkpoint):
checkpoint['example_count'] = self.example_count
checkpoint['custom_global_step'] = self.custom_global_step

def on_load_checkpoint(self, checkpoint):
self.example_count = checkpoint['example_count']
self.custom_global_step = checkpoint['custom_global_step']

def on_train_start(self):
print(f"GPU {self.global_rank}: Starting/Resuming training. Example count: {self.example_count}, Global Step: {self.custom_global_step}")

def configure_optimizers(self):
return torch.optim.Adam(self.parameters(), lr=0.02)

def main():
parser = argparse.ArgumentParser(description="Train a dummy model with a unique run name.")
parser.add_argument("run_name", type=str, help="Unique name for this training run")
args = parser.parse_args()
print("="*30, args.run_name, "="*30)

log_dir = os.path.join("logs", args.run_name)
os.makedirs(log_dir, exist_ok=True)

model = DummyModel()
dm = DummyDataModule(batch_size=4)

checkpoint_callback = ModelCheckpoint(
dirpath=os.path.join(log_dir, 'checkpoints'),
filename='model-{epoch:02d}-{train_loss:.2f}',
save_top_k=1,
verbose=True,
monitor='train_loss',
mode='min',
every_n_epochs=1,
save_last=True
)

trainer = L.Trainer(
max_epochs=100,
devices=4,
accelerator='gpu',
strategy='ddp',
callbacks=[checkpoint_callback],
plugins=[SLURMEnvironment(auto_requeue=True, requeue_signal=signal.SIGUSR1)],
default_root_dir=log_dir,
use_distributed_sampler=False,
)

trainer.fit(model, datamodule=dm, ckpt_path="last")

if __name__ == '__main__':
main()
```

dummy_slurm.sh
```bash
#!/bin/bash
#SBATCH --job-name=auto_requeue_test
#SBATCH -C h100
#SBATCH -A ycy@h100
#SBATCH --nodes=1
#SBATCH --qos=qos_gpu_h100-dev
#SBATCH --ntasks-per-node=4
#SBATCH --gpus-per-node=4
#SBATCH --time=00:3:00
#SBATCH --signal=SIGUSR1@30 # Send signal 30 seconds before time limit

# Load any necessary modules or activate your environment here
# For example:
module purge
module load arch/h100
module load pytorch-gpu/py3/2.4.0
export PYTHONUSERBASE=$WORK/python_envs/worldmodel

echo "Starting job at $(date)"

# Generate a unique run name using the current date and time
RUN_NAME="run_${SLURM_JOB_ID}"

# Run the Python script with the unique run name
srun python dummy_model.py "$RUN_NAME"

echo "Job ended or requeued at $(date)"
```

### Environment

Current environment

```
- PyTorch Lightning Version: 2.4.0
- PyTorch Version: 2.4.0
- Python version: 3.11.9
- How you installed Lightning(`conda`, `pip`, source): pip
```

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

The failure is identified in src/lightning/pytorch/trainer/connectors/checkpoint_connector.py around ckpt_path="hpc" handling; start there and trace the missing-checkpoint path. Reproduce with dummy_model.py and dummy_slurm.sh, then verify that the first run starts without an HPC checkpoint and a later requeued run resumes from the created checkpoint.

Written by the indexing model from the issue text.

Assessment

Tech stack
hpc, python
Domain
distributed-systems, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Clearly specified
Newbie friendliness
38/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.