如何使用DeepSpeed来训练大模型

news2025/6/2 14:32:16

🔥 DeepSpeed是什么？

DeepSpeed 是微软开源的一个 分布式训练加速库。
能帮助我们：
- 高效训练大模型（百亿、千亿参数规模）
- 节省显存、加速训练
- 支持 ZeRO 优化器、Offload、混合精度（FP16/BF16）、梯度累积
- 快速启动多机多卡训练

总结一句话：

DeepSpeed = 大模型训练神器，尤其适合 SFT、预训练、微调阶段。

🌈 DeepSpeed 安装

1️⃣ 安装基础依赖

通常只需要：

pip install deepspeed

对于更大规模训练，可以加上：

pip install deepspeed[all]

确保安装了 PyTorch >= 1.12。

🚀 DeepSpeed 快速上手（训练脚本改造）

2️⃣ 修改训练脚本（以 PyTorch / Hugging Face 为例）

🧩 （1）DeepSpeed CLI 启动

假设你已经有一个 train.py（PyTorch训练脚本）：

deepspeed train.py --deepspeed ds_config.json

ds_config.json：DeepSpeed配置文件（稍后详细讲）。

🧩 （2）代码适配（只需两步！）

✅ a. 导入 deepspeed：

import deepspeed

✅ b. 替换优化器 & 模型初始化：

model_engine, optimizer, _, _ = deepspeed.initialize(
    args=your_args,
    model=model,
    optimizer=optimizer,
    model_parameters=model.parameters(),
    config="ds_config.json"
)

✅ c. 训练 loop 改为：

for batch in dataloader:
    outputs = model_engine(batch)
    loss = outputs.loss
    model_engine.backward(loss)
    model_engine.step()

🎯 小结：只需 initialize 和 model_engine 替换，几行代码搞定！

🔍 DeepSpeed配置文件（`ds_config.json`）详解

这是 DeepSpeed 的核心，控制训练的优化策略。常见配置如下：

{
  "train_batch_size": 32,
  "train_micro_batch_size_per_gpu": 4,
  "gradient_accumulation_steps": 8,
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "cpu"
    },
    "offload_param": {
      "device": "cpu"
    }
  },
  "fp16": {
    "enabled": true
  },
  "gradient_clipping": 1.0,
  "steps_per_print": 100,
  "wall_clock_breakdown": false
}

⚙️ 常见配置解释：

参数	含义	推荐值 / 建议
`train_batch_size`	全局 batch size	必须设置
`train_micro_batch_size_per_gpu`	每个GPU的 batch size	看显存而定
`gradient_accumulation_steps`	梯度累积步数	`train_batch_size / (num_gpus * micro_batch_size)`
`zero_optimization`	ZeRO 优化器	stage 1/2/3
`offload_optimizer`	优化器 offload	省显存，慢一点
`offload_param`	参数 offload	stage 3 时常用
`fp16` / `bf16`	混合精度	`true`
`gradient_clipping`	梯度裁剪	1.0

📦 Hugging Face 🤗 集成 DeepSpeed

Hugging Face Transformers 已原生支持 DeepSpeed！
只需在 trainer 里加上 --deepspeed 参数即可！

✅ 步骤：
1️⃣ 准备 ds_config.json
2️⃣ 命令行运行：

accelerate config  # 配置训练
accelerate launch --multi_gpu --deepspeed ds_config.json train.py

✅ 代码示例：

from transformers import Trainer, TrainingArguments

training_args = TrainingArguments(
    output_dir="./results",
    per_device_train_batch_size=2,
    per_device_eval_batch_size=2,
    gradient_accumulation_steps=8,
    fp16=True,
    deepspeed="ds_config.json",  # 只需加这一行！
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)

trainer.train()

💡 高级技巧

✅ ZeRO-3 + Offload
最大化节省显存（即使只有 24GB 显卡也能训练 65B 模型！）
✅ Activation Checkpointing
减少显存占用，开启方式：

"activation_checkpointing": {
  "partition_activations": true,
  "contiguous_memory_optimization": true
}

✅ 梯度累积
模拟大 batch size，显存不够时的必杀技。
✅ DeepSpeed Inference Engine
支持推理加速，适合部署阶段。

🌳 项目结构

sft_project/
├── data/
│   ├── train.jsonl
│   └── val.jsonl
├── model/
│   └── (预训练模型文件夹，如LLaMA、Baichuan)
├── deepspeed_config/
│   └── ds_config.json
├── train.py
├── requirements.txt
└── README.md

🎓 总结

你想做什么？	如何用DeepSpeed？
训练大模型	用 `deepspeed` 启动，写好 `ds_config.json`
不想改代码	Hugging Face `Trainer` + `--deepspeed` 参数
显存不够	开启 ZeRO-3 + Offload + FP16/BF16
多机多卡训练	`deepspeed --num_gpus=8` 或 `accelerate launch`
部署	DeepSpeed Inference 加速推理