SLURM集群训练遇时间限制终止时如何保存PyTorch Lightning模型权重

1. 调整SLURM作业提交命令 #

首先配置SLURM在作业超时前提前发送告警信号，预留出模型保存的时间。以下示例为超时前300秒（5分钟）发送SIGUSR1信号，提交命令修改为：

sbatch --gpus=1 -t 100 --signal=SIGUSR1@300 python train.py

你可以根据自己模型保存的耗时调整预留时长，确保足够完成权重写入操作。

2. 配置PyTorch Lightning Trainer #

PyTorch Lightning 1.5及以上版本默认内置SLURM信号处理能力，仅需配置对应参数即可实现收到超时信号后自动保存权重再优雅退出，配置示例如下：

from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import ModelCheckpoint

# 配置Checkpoint回调
checkpoint_callback = ModelCheckpoint(
    save_last=True, # 始终保留最新的权重文件
    every_n_train_steps=1000, # 可自定义每N个训练步保存一次，进一步降低进度损失风险，步长可根据实际情况调整
    dirpath="./checkpoints/",
    filename="model-{step}-{loss:.2f}"
)

# 初始化Trainer
trainer = Trainer(
    callbacks=[checkpoint_callback],
    enable_slurm_signal_handlers=True, # 开启SLURM信号监听，收到提前告警信号后自动触发权重保存
    # 其余原有配置保持不变，比如 accelerator="gpu", devices=1 等
)

3. 恢复训练 #

后续恢复训练时，仅需在Trainer初始化时传入保存好的最新权重路径即可接着之前的进度继续训练，无需等待完整epoch跑完：

trainer = Trainer(
    ckpt_path="./checkpoints/last.ckpt",
    # 其余配置保持不变
)

内容的提问来源于stack exchange，提问作者Kenenbek Arzymatov