训练容错
训练过程中的错误和故障可能会中断整个训练作业,导致丢失训练状态。AIStudio 训练服务模块提供平台级别的容错能力,可以自动恢复训练(训练脚本必须支持优雅重启),避免重新开始。
接入 AIStudio 的训练容错能力需要满足以下要求:
- 为任务开启「容错」功能。
- 确保训练脚本支持优雅重启。您的训练脚本应定期保存 Checkpoint,并在开始训练时加载上次保存的 Checkpoint。
了解容错流程
- Bootcheck(启动检测):在任务启动时自动检查关键的软件和硬件环境(GPU/RDMA/IO/通信)。仅在发现掉卡错误时会阻断任务并重新调度健康节点。其他情况下会报告并继续执行用户代码。
- Command(模型训练):执行用户代码。
- Troubleshoot(异常定位):在用户代码执行失败时,检测硬件健康情况。仅在发现掉卡错误时会重新调度健康节点。其他情况下会报告并继续执行用户代码。如果用户代码有问题,可能将反复循环,直到达到最大容错次数。
TIP
AIStudio 任务模块提供 atlctl
命令行调试工具,支持用户从 Web Terminal 登录任务 Worker 对任务执行环境进行原地调试。上图未体现生命周期中的 atlctl 调试状态。详见 atlctl 调试工具。
为任务开启容错功能
在创建任务时,请确保「容错」开关为开启状态。
如遇任务异常,平台会自动重启 Worker 或进程。默认最多重启 1 次,最大 10 次。如果超过最大重启次数,任务将失败并结束。
训练脚本支持优雅重启
支持优雅重启指每个进程的入口点首先加载并初始化上次保存的 Checkpoint,然后继续训练。因此,在任何故障情况下,您只会丢失自上次保存快照以来的训练进度。
为支持优化重启,训练脚本中应存在 load_snapshot
、save_snapshot(snapshot_path)
等类似逻辑。
python
def main():
load_snapshot(snapshot_path)
initialize()
train()
def train():
for batch in iter(dataset):
train_step(batch)
if should_checkpoint:
save_snapshot(snapshot_path)
- 具体实现可参考 Pytorch 官方文档 Saving and loading snapshots - Fault-tolerant Distributed Training with torchrun
- 使用 DDP 时,一个优化方法是只在一个进程中保存模型,然后在所有进程中加载,从而减少写入开销。参见 Pytorch 官方文档的详细示例:Save and Load Checkpoints - Getting Started with Distributed Data Parallel
追踪容错日志
开启容器功能后,可在任务的详情页中找到容错日志标签。其中记录了每一轮容错时开机检测结果、异常定位、重启操作等记录。
任务详情的容错日志会提示当前任务的最大重启次数、已重启次数、因硬件问题导致的重调度次数。
启动检测(Bootcheck)与训练失败时执行的异常定位(Troubleshoot),均会检查关键的软件和硬件环境(GPU/RDMA/IO/通信)。正常输出示例:
shell
2025-01-08 15:54:15 Bootcheck starting...
2025-01-08 15:54:16 [IO Check] IO Path: /mnt/public, Read 7.6 GB/s, Write 408 MB/s.
2025-01-08 15:54:18 [RDMA Check] NIC Type: nil; Total NICs: 0, NotReady: 0.
2025-01-08 15:54:25 [Runtime Check] PyTorch: v2.3.0a0+40ec155e58.nv24.03; CUDA: 12.4; NCCL: (2, 20, 5).
2025-01-08 15:54:25 [GPU Check] GPU Type: NVIDIA A100-PCIE-40GB; Total GPUs: 2; Not Ready: 0
2025-01-08 15:54:31 [Communication Check] AllReduce connection passed.
2025-01-08 15:54:31 Check Finished.
NOTE
任务详情页的任务日志也会输出容错日志。