设置训练故障优雅退出

使用场景

随着模型规模和数据集的急剧增长，需要利用大规模的训练集训练大规模的神经网络。在大规模集群分布式训练时，会遇到集群中某个芯片、某台服务器故障，导致分布式训练任务失败。优雅退出是指中断的训练任务支持自动恢复，并可以在上一次训练中断的基础上继续训练，而不用从头开始。

约束限制

特性原理

特性使用操作

1. 安装优雅退出二进制包

   通过ma_pre_start.sh安装whl包。

   echo "[ma-pre-start] Enter the input directory"
   cd /home/ma-user/modelarts/inputs/data_url_0/​
   echo "[ma-pre-start] Start to install mindx-elastic 0.0.1版本"
   export PATH=/home/ma-user/anaconda/bin:$PATH
   pip install ./mindx_elastic-0.0.1-py3-none-any.whl
   echo "[ma-pre-start] Clean run package"
   sudo rm -rf ./script ./*.run ./run_package *.whl
   echo "[ma-pre-start] Set ENV"
   export GLOG_v=2    # 当前使用诊断模式需要用户手动设置成INFO日志级别 echo "[ma-pre-start] End"

2. 创建训练任务
   • 约束：MindSpore版本要求1.6.0及以上。
   • 修改样例代码，增加如下内容：

     # 载入依赖接口
     from mindx_elastic.terminating_message import ExceptionCheckpoint
     ...
     
     if args_opt.do_train:
     dataset = create_dataset()
     loss_cb = LossMonitor()
     cb = [loss_cb]
     if int(os.getenv('RANK_ID')) == 0:
     batch_num = dataset.get_dataset_size()
     # 开启优雅退出保存
     config_ck = CheckpointConfig(save_checkpoint_steps=batch_num,
     keep_checkpoint_max=35,
     async_save=True,
     append_info=[{"epoch_num": cur_epoch_num}],
     exception_save=True)
     
     ckpoint_cb = ModelCheckpoint(prefix="train_resnet_cifar10",
     directory=args_opt.train_url,
     config=config_ck)
     # 定义优雅退出ckpt保存callback
     ckpoint_exp = ExceptionCheckpoint(
     prefix="train_resnet_cifar10",
     directory=args_opt.train_url,
     config=config_ck)
     # 添加优雅退出ckpt保存callback
     cb += [ckpoint_cb, ckpoint_exp]