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

[F-Barto]

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

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

#!/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

[arijit-hub]

Hi @F-Barto,
You do not need to specify ckpt_path="hpc". In the current setting of lightning (url) it always searches for the "hpc" path internally first, then if it doesn't find it, it takes the ckpt_path that you specify in your trainer.fit(..., ckpt_path=ckpt_path). So you shouldn't specify ckpt_path="hpc" and instead just do something like this:

## Option:1
## If you want to have an option for manual resuming
## Have a flag for resuming (args.resume)
# This will do autorequeue (or) resume with last (or) run from start.
trainer.fit(..., ckpt_path="/path/to/saved/checkpoints/last.ckpt" if args.resume else None) 

## Option: 2
## No manual resuming, just autoreque (or) run from start
trainer.fit(..., ckpt_path=None) 

Typically what happens is when you run a .sh file and your job is about to hit the wall-time, lightning automatically creates a temporary checkpoint (hpc_ckpt_*.ckpt) in the default_root_dir set by the user in trainer. Then when the job restarts, lightning automatically searches the default_root_dir folder, and if the hpc_ckpt_*.ckpt file is there, it would load it and resume training.

I hope this helps and it was clear enough. Let me know if there is something which is still confusing.