MT-DeepSpeed 入门

MT-DeepSpeed 简介

环境准备

MT-DeepSpeed 是面向摩尔线程 MUSA 平台适配的分布式训练加速组件，提供 ZeRO 优化、梯度累积、混合精度和多卡通信等能力。它通常与 Transformers、PyTorch 或上层训练脚本配合使用。

建议使用摩尔线程官方训练镜像。当前示例使用 MTT S5000、MUSA Driver 4.3.7 和 training-suite:v2.1.5-musa-4.3.7。

创建并启动容器

sudo docker create --privileged --env MTHREADS_VISIBLE_DEVICES=all --net host -v /data/:/data/ --name deepspeed registry.mthreads.com/mcctest/training-suite:v2.1.5-musa-4.3.7 sleep infinity
sudo docker start deepspeed
sudo docker exec -it deepspeed bash
service ssh restart

验证 GPU 和 DeepSpeed

mthreads-gmi
# ./mccl_test

使用 ds_report 确认输出包含 ds_accelerator to musa 和 musa_fused_adam。

使用 MT-DeepSpeed 进行模型微调

1. 准备模型和数据

pip install modelscope
modelscope download --model Qwen/Qwen3-4B --local_dir /data/models/Qwen3-4B

训练数据应按照上层训练脚本要求准备，对话式 SFT 数据通常包含 messages、role 和 content 字段。

2. 配置 ZeRO 和混合精度

DeepSpeed 配置通过 JSON 文件传入训练脚本。典型的 BF16 + ZeRO-2 配置如下：

{
  "bf16": {"enabled": true},
  "fp16": {"enabled": false},
  "zero_optimization": {"stage": 2, "overlap_comm": true, "reduce_scatter": true, "contiguous_gradients": true},
  "gradient_accumulation_steps": "auto",
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto"
}

显存更紧张或模型规模更大时，可以评估 ZeRO-3，并同步检查通信开销、保存方式和显存占用。

3. 启动多卡训练

deepspeed --num_gpus=4 train_qwen3_sft.py \
  --model_path /data/models/Qwen3-4B \
  --data_path /data/data/train.json \
  --output_dir ./output/qwen3-4b-sft-ds \
  --deepspeed_config ./configs/ds_config_qwen3_4b_zero2.json \
  --bf16 --per_device_batch_size 2 --gradient_accumulation_steps 8

其中 --num_gpus 指定 GPU 数量，--deepspeed_config 指定 ZeRO、精度和 batch 配置。多机训练还需要配置主节点地址、端口、节点数和 SSH 连通性。

4. 监控和排障

• 使用 mthreads-gmi 检查各卡显存与利用率；
• 遇到 OOM、MCCL 通信错误时，优先检查 ZeRO stage、batch size、梯度累积和 GPU 数量；
• 训练结束后确认 checkpoint 已生成且分布式进程退出。