🔥 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_optimizationZeRO 优化器stage 1/2/3
offload_optimizer优化器 offload省显存,慢一点
offload_param参数 offloadstage 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 加速推理

Logo

北京人形旗下天工造物具身智能开源社区,聚焦具身天工与慧思开物两大平台

更多推荐