xMIG (xPU 迁移技术)


一  整体介绍:

alt png     xMig 是一个为大规模语言模型训练/推理设计的轻量级、快速且透明的错误恢复框架,采用即时(Just-In-Time)方式进行错误处理。其核心目标是通过高效的机制减少训练过程中的错误恢复开销,同时保持系统的稳定性和性能。

核心特点

  • 透明性与无周期检查点:Mnemosyne 继承了微软即时检查点技术的优势,无需周期性保存检查点,从而降低了训练过程中的性能开销;
  • 定制化错误恢复机制:通过优化的恢复流程,最大限度地减少日常训练中的额外开销;
  • 支持运行时调整通信组:提供灵活的集合通信CC组调整功能,适应动态训练需求;

技术实现

  • 设备代理进行故障监控和故障拦截 xPU 相关的错误会通过API进行钩子捕获,检测到错误后,设备代理停止向上层框架继续传播错误,进入故障恢复状态;

  • 故障恢复工作流 对于可重启恢复的软件/驱动错误,进行原地恢复,重启故障节点并进行日志重放、CCL 重建、DP 参数拉取;对于不可恢复的错误(如硬件故障),使用CRIU保存CPU进程状态,迁移CPU状态到新节点,随后恢复GPU状态恢复时,使用 Flexible CCL 加速集合通信重建;

  • 恢复完成后,设备代理正确响应上层框架;

二  安装教程:

  1. 设置环境变量
$ export CUDA_HOME=\<cuda-path\>
$ export NCCL_HOME=\<xmig-path\>/GPU/target/nccl-2.27.7-1
$ export PATH=$CUDA_HOME/bin:$PATH
$ export LD_LIBRARY_PATH=$NCCL_HOME/build/lib:$CUDA_HOME/lib64:$LD_LIBRARY_PATH
  1. 编译 NCCL 和 xMIG
$ cd <xmig-path>/GPU/
$ make
  1. 运行测试脚本
$ cd <xmig-path>/test/
# 指定run.sh脚本使用的GPU设备和数量,默认配置如下:
# CUDA_VISIBLE_DEVICES=0,1,2,3
# torchrun --nproc_per_node=4
$ bash run.sh

Docker容器内编译

我们提供了预构建的Docker镜像,包含xMIG及其依赖项,方便用户快速部署和测试。请按照以下步骤使用Docker容器:

  1. 编译Docker镜像
$ cd <host-xmig-path>/
$ docker build -t <docker-image-name> .

如果需要配置代理,将证书复制到xMIG根目录下,并使用以下命令构建镜像:

$ cp <crt-file-path>/<crt-file-name> <host-xmig-path>/
$ docker build \
    --build-arg http_proxy=<http-proxy> \
    --build-arg https_proxy=<https-proxy> \
    --build-arg CA_PATH=<crt-file-name> \
    -t <docker-image-name> .
  1. 运行Docker容器
$ docker run -it --rm --gpus all --ipc=host <docker-image-name> /bin/bash
  1. 在Docker容器内编译xMIG
$ cd /workspace/xmig/GPU/
$ make
  1. 运行测试脚本
$ cd /workspace/xmig/GPU/test/
$ bash run.sh # 配置CUDA_VISIBLE_DEVICES和torchrun参数

三  使用场景:

    适用于需要高效错误恢复的大规模分布式训练任务,尤其是在训练大规模语言模型时,能够显著减少因错误导致的时间和资源浪费:

  • 原地恢复:通过集合通信修复,原地就行修复;
  • 容器迁移:通过CRIU进行容器迁移,进行容器迁移修复;

四  参考链接:

  1. Mnemosyne: Lightweight and Fast Error Recovery for LLM Training in a Just-In-Time Manner
  2. PhoenixOS: Concurrent OS-level GPU Checkpoint and Restore with Validated Speculation
  3. Just-In-Time Checkpointing: Low Cost Error Recovery from Deep Learning Training Failures
  4. Singularity: Planet-Scale, Preemptive and Elastic Scheduling of AI Workloads