当程序在 SLURM 集群上学习时收到 TIME LIMIT 时保存模型权重

Question

我使用用 pytorch_lightning (pytorch) 编写的深度学习模型，并在 slurm 集群上训练它们。我这样提交作业：

sbatch --gpus=1 -t 100 python train.py

当请求 GPU 时间结束时，slurm 杀死我的程序并显示这样的消息：

Epoch 0: : 339it [01:10,  4.84it/s, loss=-34]  slurmstepd: error: *** JOB 375083 ON cn-007 CANCELLED AT 2021-10-04T22:20:54 DUE TO TIME LIMIT ***

如何配置 Trainer 在可用时间结束时保存模型？

我知道每个纪元后 automatic saving，但我只有一个持续时间 >10 小时的长纪元，所以这种情况对我不起作用。

Answer 1

您可以使用 Slurm 的信号机制在时间限制的特定秒数内将信号传递给您的应用程序（请参阅 man sbatch）。在您的提交脚本中，使用 --signal=USR1@30 在达到时限前 30 秒发送 USR1。您的提交脚本将包含这些行：

#SBATCH -t 100
#SBATCH --signal=USR1@30
srun python train.py

然后，在您的代码中，您可以像这样处理该信号：

import signal

def handler(signum, frame):
    print('Signal handler got signal ', signum)
    # e.g. exit(0), or call your pytorch save routines

# enable the handler
signal.signal(signal.SIGUSR1, handler)

# your code here

您需要通过 srun 调用您的 Python 应用程序，以便 Slurm 能够将信号传播到 Python 进程。（您可以在命令行上使用 --signal 到 sbatch，我倾向于编写独立的提交脚本 :)）

编辑：This link 很好地总结了信号传播和 Slurm 所涉及的问题。

当程序在 SLURM 集群上学习时收到 TIME LIMIT 时保存模型权重

Save a model weights when a program receives TIME LIMIT while learning on a SLURM cluster

cluster-computing

slurm

deep-learning

pytorch

pytorch-lightning