24 DeepSpeed与ZeRO优化

DeepSpeed 简介

DeepSpeed是由Microsoft开发的大规模分布式训练优化框架，旨在支持更大规模的模型训练，提供更多显存优化策略。

主要优势和功能：

3D并行支持：灵活组合数据并行、流水线并行和张量并行，支持万亿参数模型训练
ZeRO优化：零冗余优化器，大幅减少显存占用
ZeRO-Offload：将部分计算offload到CPU/NVMe，单卡也能训练百亿参数模型
Sparse Attention：支持更长序列，速度提升6倍
1-bit Adam：减少5倍通信量，通信受限场景下训练提速3.5倍

主流训练方案：GPU + PyTorch + Megatron-LM + DeepSpeed 是目前业界训练超大模型的标准组合。

ZeRO-1/2/3 原理

ZeRO核心思想

ZeRO（Zero Redundancy Optimizer）去除数据并行中的冗余存储，每张卡只存储一部分模型状态，从而大幅减少显存占用。在执行逻辑上它仍是数据并行，但能达到模型并行的显存优化效果。

模型训练时，显存主要被两部分占用：

模型状态：参数、梯度、优化器状态（优化器状态占比约75%）
剩余状态：激活值、临时缓冲区、显存碎片

以1.5B参数GPT-2为例，fp16混合精度下，虽然参数只需要3GB，但模型状态总共需要24GB，优化器状态是显存最大消耗者。

ZeRO三个阶段

ZeRO分三个阶段，逐步对不同部分进行分片存储：

点击放大查看

ZeRO阶段	优化器状态	梯度	参数	显存减少	通信量增加
stage 0	完整存储	完整存储	完整存储	1x	基准
stage 1	分片存储	完整存储	完整存储	4x	不变
stage 2	分片存储	分片存储	完整存储	8x	不变
stage 3	分片存储	分片存储	分片存储	N倍 (N=GPU数)	+50%

举例：64块GPU，ZeRO-3可以把显存减少到1/64，让千亿模型训练成为可能。

点击放大查看

ZeRO-Offload / ZeRO-Infinity

ZeRO-Offload原理

ZeRO-Offload将模型状态卸载（offload）到CPU内存，利用CPU内存来扩展可训练模型大小，让GPU单卡也能训练更大的模型。

计算流程：

点击放大查看

关键优化：计算与通信重叠。GPU在反向传播阶段，一边计算梯度，一边把已经算好的梯度bucket传输给CPU，当反向结束，CPU已经拿到大部分梯度了。

ZeRO-Infinity

ZeRO-Infinity进一步扩展了offload思想，可以将参数offload到NVMe固态硬盘，理论上可以用有限GPU显存训练任意大的模型：

offload到CPU：利用CPU内存
offload到NVMe：利用SSD存储空间，进一步扩展

性能特点

从左到右，速度越来越慢，显存需求越来越小：

速度：Stage 0 (DDP) > Stage 1 > Stage 2 > Stage 2 + offload > Stage 3 > Stage 3 + offload
显存：Stage 0 (DDP) < Stage 1 < Stage 2 < Stage 2 + offload < Stage 3 < Stage 3 + offload

DeepSpeed 配置与使用

安装

bash

pip install deepspeed==0.8.1
sudo apt-get install openmpi-bin libopenmpi-dev
pip install mpi4py

基本使用流程

导入并初始化DeepSpeed引擎

python

import deepspeed

# 定义模型、优化器
model = Model()
criterion = nn.CrossEntropyLoss()
optimizer = build_optimizer(model, args)

# DeepSpeed配置
deepspeed_config = {
    "train_micro_batch_size_per_gpu": 32,
    "gradient_accumulation_steps": 1,
    "optimizer": {
        "type": "AdamW",
        "params": {"lr": 3e-5}
    },
    "fp16": {"enabled": True},
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": true
        }
    }
}

# 初始化DeepSpeed引擎
model_engine, optimizer, _, _ = deepspeed.initialize(
    config=deepspeed_config,
    model=model,
    model_parameters=model.parameters()
)

训练循环

python

for batch in dataloader:
    # 前向
    outputs = model_engine(inputs)
    loss = criterion(outputs, labels)
    
    # 反向
    model_engine.backward(loss)
    
    # 步进
    model_engine.step()

保存和加载

python

# 保存
model_engine.save_checkpoint(save_path, save_latest=True)

# 加载
model_engine.load_checkpoint(load_path, load_module_only=True)

# 转换为完整pytorch模型
# 使用zero_to_fp32.py脚本
python zero_to_fp32.py . pytorch_model.bin

运行

bash

deepspeed --hostfile=hostfile train.py --deepspeed_config config.json

重要配置参数说明

ZeRO-2 配置示例：

json

{
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "contiguous_gradients": true
  }
}

ZeRO-3 配置示例：

json

{
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "offload_param": {
      "device": "cpu",
      "pin_memory": true
    },
    "overlap_comm": true,
    "contiguous_gradients": true,
    "sub_group_size": 1e9,
    "reduce_bucket_size": "auto",
    "stage3_prefetch_bucket_size": "auto",
    "stage3_param_persistence_threshold": "auto",
    "stage3_max_live_parameters": 1e9,
    "stage3_max_reuse_distance": 1e9,
    "stage3_gather_16bit_weights_on_model_save": true
  }
}

关键参数说明：

overlap_comm: 是否让通信和计算重叠，开启可加速
allgather_bucket_size / reduce_bucket_size: 控制分桶大小，越大通信越快但占更多显存。OOM时需要调小
stage3_max_live_parameters: 控制GPU上同时保留的完整参数数量上限，OOM时调小

ZeRO 各阶段显存占用对比

点击放大查看

DeepSpeed 调参指南

OOM调参步骤（按顺序尝试）

减小batch size，使用梯度累积：batch_size=1 + 梯度累积
开启梯度检查点：gradient_checkpointing
尝试 ZeRO stage 2
尝试 ZeRO stage 2 + offload optimizer 到CPU
尝试 ZeRO stage 3
尝试 offload param 到CPU
如果还是OOM，尝试 offload 到NVMe (ZeRO-Infinity)
减小allgather_bucket_size和reduce_bucket_size

估算ZeRO-3显存需求

DeepSpeed提供显存估算工具：

python

from transformers import AutoModel
from deepspeed.runtime.zero.stage3 import estimate_zero3_model_states_mem_needs_all_live

model = AutoModel.from_pretrained("bigscience/T0_3B")
estimate_zero3_model_states_mem_needs_all_live(model, num_gpus_per_node=2, num_nodes=1)

输出示例：

per CPU  |  per GPU |   Options
 70.00GB |   0.25GB | offload_param=cpu , offload_optimizer=cpu
 62.23GB |   2.84GB | offload_param=none, offload_optimizer=cpu
  0.74GB |  23.58GB | offload_param=none, offload_optimizer=none

可以根据这个输出选择合适的配置。

常见问题排查

ModuleNotFoundError: No module named 'torch._six'
- 这是新版本PyTorch兼容性问题，注释掉import，手动定义：
  python
  1# from torch._six import string_classes 2int_classes = int 3string_classes = str
loss是NaN
- 检查：如果你用的是bf16预训练模型，不要用fp16训练，改用bf16或fp32
启动时进程被杀死，无traceback
- 这是GPU显存不够，需要加卡或者更激进的offload
单卡也能用DeepSpeed吗？
- 可以。使用ZeRO-offload把数据放到CPU，降低GPU显存需求，还能减少显存碎片
ZeRO-3比ZeRO-2慢很多，怎么优化？
- 增大stage3_param_persistence_threshold
- 关闭offload_params（能显著提升性能）

多机配置要点

ssh配置：所有节点互信免密登录
环境一致：代码路径、python路径、库版本都要一致
hostfile：列出所有节点和可用GPU数
model1 slots=8 model2 slots=8

面试常见问题

Q1: ZeRO和传统数据并行的区别是什么？

A: 传统数据并行中，每个GPU都保存完整的参数、梯度、优化器状态，有大量冗余存储，显存利用率低。ZeRO通过分片存储，每个GPU只保存一部分，去除冗余，在数据并行的并行度下，达到了模型并行的显存缩减效果。

Q2: ZeRO的三个阶段分别做了什么？

ZeRO-1：只分片优化器状态，显存减少4倍
ZeRO-2：分片优化器状态 + 梯度，显存减少8倍
ZeRO-3：分片优化器状态 + 梯度 + 参数，显存减少和GPU数线性相关

Q3: ZeRO-Offload的原理是什么？

A: ZeRO-Offload把优化器状态和参数放到CPU内存，只在需要时传到GPU计算，完成后再传回去。用CPU内存换GPU显存，让小GPU也能训练大模型。关键优化是计算和通信重叠，隐藏传输延迟。

Q4: DeepSpeed Accelerate 有什么区别？

DeepSpeed：面向超大模型训练，提供ZeRO、offload、3D并行等高级显存优化，功能强大但配置复杂
Accelerate：HuggingFace推出，简化分布式训练，API简单易用，封装了DDP、混合精度，适合中小型模型分布式训练

Q5: 怎么选择ZeRO stage？

A: 根据GPU显存从多到少依次尝试：

显存充足：stage 0（等价DDP）最快
显存够但有点紧：stage 1
显存比较紧张：stage 2
显存非常紧张：stage 2 + offload
模型非常大：stage 3
模型极大：stage 3 + offload

Q6: ZeRO和3D并行是什么关系，为什么有了ZeRO还需要3D并行？

A: ZeRO是显存优化的数据并行，3D并行是数据+张量+流水线的组合并行。根据论文，当GPU数量增加到千量级，3D并行的效率比纯ZeRO3更高。ZeRO在大规模集群上通信效率不如3D并行。实际中常常结合使用：Megatron-LM做张量+流水线并行，DeepSpeed做ZeRO数据并行。

24 DeepSpeed与ZeRO优化

DeepSpeed 简介

DeepSpeed是由Microsoft开发的大规模分布式训练优化框架，旨在支持更大规模的模型训练，提供更多显存优化策略。

主要优势和功能：

3D并行支持：灵活组合数据并行、流水线并行和张量并行，支持万亿参数模型训练
ZeRO优化：零冗余优化器，大幅减少显存占用
ZeRO-Offload：将部分计算offload到CPU/NVMe，单卡也能训练百亿参数模型
Sparse Attention：支持更长序列，速度提升6倍
1-bit Adam：减少5倍通信量，通信受限场景下训练提速3.5倍

主流训练方案：GPU + PyTorch + Megatron-LM + DeepSpeed 是目前业界训练超大模型的标准组合。

ZeRO-1/2/3 原理

ZeRO核心思想

模型训练时，显存主要被两部分占用：

模型状态：参数、梯度、优化器状态（优化器状态占比约75%）
剩余状态：激活值、临时缓冲区、显存碎片

以1.5B参数GPT-2为例，fp16混合精度下，虽然参数只需要3GB，但模型状态总共需要24GB，优化器状态是显存最大消耗者。

ZeRO三个阶段

ZeRO分三个阶段，逐步对不同部分进行分片存储：

点击放大查看

ZeRO阶段	优化器状态	梯度	参数	显存减少	通信量增加
stage 0	完整存储	完整存储	完整存储	1x	基准
stage 1	分片存储	完整存储	完整存储	4x	不变
stage 2	分片存储	分片存储	完整存储	8x	不变
stage 3	分片存储	分片存储	分片存储	N倍 (N=GPU数)	+50%

举例：64块GPU，ZeRO-3可以把显存减少到1/64，让千亿模型训练成为可能。

点击放大查看

ZeRO-Offload / ZeRO-Infinity

ZeRO-Offload原理

ZeRO-Offload将模型状态卸载（offload）到CPU内存，利用CPU内存来扩展可训练模型大小，让GPU单卡也能训练更大的模型。

计算流程：

点击放大查看

关键优化：计算与通信重叠。GPU在反向传播阶段，一边计算梯度，一边把已经算好的梯度bucket传输给CPU，当反向结束，CPU已经拿到大部分梯度了。

ZeRO-Infinity

ZeRO-Infinity进一步扩展了offload思想，可以将参数offload到NVMe固态硬盘，理论上可以用有限GPU显存训练任意大的模型：

offload到CPU：利用CPU内存
offload到NVMe：利用SSD存储空间，进一步扩展

性能特点

从左到右，速度越来越慢，显存需求越来越小：

速度：Stage 0 (DDP) > Stage 1 > Stage 2 > Stage 2 + offload > Stage 3 > Stage 3 + offload
显存：Stage 0 (DDP) < Stage 1 < Stage 2 < Stage 2 + offload < Stage 3 < Stage 3 + offload

DeepSpeed 配置与使用

安装

bash

pip install deepspeed==0.8.1
sudo apt-get install openmpi-bin libopenmpi-dev
pip install mpi4py

基本使用流程

导入并初始化DeepSpeed引擎

python

import deepspeed

# 定义模型、优化器
model = Model()
criterion = nn.CrossEntropyLoss()
optimizer = build_optimizer(model, args)

# DeepSpeed配置
deepspeed_config = {
    "train_micro_batch_size_per_gpu": 32,
    "gradient_accumulation_steps": 1,
    "optimizer": {
        "type": "AdamW",
        "params": {"lr": 3e-5}
    },
    "fp16": {"enabled": True},
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": true
        }
    }
}

# 初始化DeepSpeed引擎
model_engine, optimizer, _, _ = deepspeed.initialize(
    config=deepspeed_config,
    model=model,
    model_parameters=model.parameters()
)

训练循环

python

for batch in dataloader:
    # 前向
    outputs = model_engine(inputs)
    loss = criterion(outputs, labels)
    
    # 反向
    model_engine.backward(loss)
    
    # 步进
    model_engine.step()

保存和加载

python

# 保存
model_engine.save_checkpoint(save_path, save_latest=True)

# 加载
model_engine.load_checkpoint(load_path, load_module_only=True)

# 转换为完整pytorch模型
# 使用zero_to_fp32.py脚本
python zero_to_fp32.py . pytorch_model.bin

运行

bash

deepspeed --hostfile=hostfile train.py --deepspeed_config config.json

重要配置参数说明

ZeRO-2 配置示例：

json

{
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "contiguous_gradients": true
  }
}

ZeRO-3 配置示例：

json

{
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "offload_param": {
      "device": "cpu",
      "pin_memory": true
    },
    "overlap_comm": true,
    "contiguous_gradients": true,
    "sub_group_size": 1e9,
    "reduce_bucket_size": "auto",
    "stage3_prefetch_bucket_size": "auto",
    "stage3_param_persistence_threshold": "auto",
    "stage3_max_live_parameters": 1e9,
    "stage3_max_reuse_distance": 1e9,
    "stage3_gather_16bit_weights_on_model_save": true
  }
}

关键参数说明：

overlap_comm: 是否让通信和计算重叠，开启可加速
allgather_bucket_size / reduce_bucket_size: 控制分桶大小，越大通信越快但占更多显存。OOM时需要调小
stage3_max_live_parameters: 控制GPU上同时保留的完整参数数量上限，OOM时调小

ZeRO 各阶段显存占用对比

点击放大查看

DeepSpeed 调参指南

OOM调参步骤（按顺序尝试）

减小batch size，使用梯度累积：batch_size=1 + 梯度累积
开启梯度检查点：gradient_checkpointing
尝试 ZeRO stage 2
尝试 ZeRO stage 2 + offload optimizer 到CPU
尝试 ZeRO stage 3
尝试 offload param 到CPU
如果还是OOM，尝试 offload 到NVMe (ZeRO-Infinity)
减小allgather_bucket_size和reduce_bucket_size

估算ZeRO-3显存需求

DeepSpeed提供显存估算工具：

python

from transformers import AutoModel
from deepspeed.runtime.zero.stage3 import estimate_zero3_model_states_mem_needs_all_live

model = AutoModel.from_pretrained("bigscience/T0_3B")
estimate_zero3_model_states_mem_needs_all_live(model, num_gpus_per_node=2, num_nodes=1)

输出示例：

per CPU  |  per GPU |   Options
 70.00GB |   0.25GB | offload_param=cpu , offload_optimizer=cpu
 62.23GB |   2.84GB | offload_param=none, offload_optimizer=cpu
  0.74GB |  23.58GB | offload_param=none, offload_optimizer=none

可以根据这个输出选择合适的配置。

常见问题排查

ModuleNotFoundError: No module named 'torch._six'
- 这是新版本PyTorch兼容性问题，注释掉import，手动定义：
  python
  1# from torch._six import string_classes 2int_classes = int 3string_classes = str
loss是NaN
- 检查：如果你用的是bf16预训练模型，不要用fp16训练，改用bf16或fp32
启动时进程被杀死，无traceback
- 这是GPU显存不够，需要加卡或者更激进的offload
单卡也能用DeepSpeed吗？
- 可以。使用ZeRO-offload把数据放到CPU，降低GPU显存需求，还能减少显存碎片
ZeRO-3比ZeRO-2慢很多，怎么优化？
- 增大stage3_param_persistence_threshold
- 关闭offload_params（能显著提升性能）

多机配置要点

ssh配置：所有节点互信免密登录
环境一致：代码路径、python路径、库版本都要一致
hostfile：列出所有节点和可用GPU数
model1 slots=8 model2 slots=8

面试常见问题

Q1: ZeRO和传统数据并行的区别是什么？

Q2: ZeRO的三个阶段分别做了什么？

ZeRO-1：只分片优化器状态，显存减少4倍
ZeRO-2：分片优化器状态 + 梯度，显存减少8倍
ZeRO-3：分片优化器状态 + 梯度 + 参数，显存减少和GPU数线性相关

Q3: ZeRO-Offload的原理是什么？

Q4: DeepSpeed Accelerate 有什么区别？

DeepSpeed：面向超大模型训练，提供ZeRO、offload、3D并行等高级显存优化，功能强大但配置复杂
Accelerate：HuggingFace推出，简化分布式训练，API简单易用，封装了DDP、混合精度，适合中小型模型分布式训练

Q5: 怎么选择ZeRO stage？

A: 根据GPU显存从多到少依次尝试：

显存充足：stage 0（等价DDP）最快
显存够但有点紧：stage 1
显存比较紧张：stage 2
显存非常紧张：stage 2 + offload
模型非常大：stage 3
模型极大：stage 3 + offload

大模型面试

大模型面试

24 DeepSpeed与ZeRO优化

24 DeepSpeed与ZeRO优化

DeepSpeed 简介

ZeRO-1/2/3 原理

ZeRO核心思想

ZeRO三个阶段

ZeRO-Offload / ZeRO-Infinity

ZeRO-Offload原理

ZeRO-Infinity

性能特点

DeepSpeed 配置与使用

安装

基本使用流程

重要配置参数说明

ZeRO 各阶段显存占用对比

DeepSpeed 调参指南

OOM调参步骤（按顺序尝试）

估算ZeRO-3显存需求

常见问题排查

多机配置要点

面试常见问题

Q1: ZeRO和传统数据并行的区别是什么？

Q2: ZeRO的三个阶段分别做了什么？

Q3: ZeRO-Offload的原理是什么？

Q4: DeepSpeed Accelerate 有什么区别？

Q5: 怎么选择ZeRO stage？

Q6: ZeRO和3D并行是什么关系，为什么有了ZeRO还需要3D并行？

24 DeepSpeed与ZeRO优化

24 DeepSpeed与ZeRO优化

DeepSpeed 简介

ZeRO-1/2/3 原理

ZeRO核心思想

ZeRO三个阶段

ZeRO-Offload / ZeRO-Infinity

ZeRO-Offload原理

ZeRO-Infinity

性能特点

DeepSpeed 配置与使用

安装

基本使用流程

重要配置参数说明

ZeRO 各阶段显存占用对比

DeepSpeed 调参指南

OOM调参步骤（按顺序尝试）

估算ZeRO-3显存需求

常见问题排查

多机配置要点

面试常见问题

Q1: ZeRO和传统数据并行的区别是什么？

Q2: ZeRO的三个阶段分别做了什么？

Q3: ZeRO-Offload的原理是什么？

Q4: DeepSpeed Accelerate 有什么区别？

Q5: 怎么选择ZeRO stage？

Q6: ZeRO和3D并行是什么关系，为什么有了ZeRO还需要3D并行？