DeepSpeed ZeRO 优化技术详解：从原理到实战，掌握大模型分布式训练的核心

引言：大模型训练的内存困境

随着GPT-4、Llama 3、Qwen等大语言模型（LLM）的参数规模突破千亿甚至万亿级别，单张GPU的显存早已无法容纳完整的模型训练。以Llama 3 70B为例，仅模型权重（FP16）就需要约140GB显存，而优化器状态（Adam的momentum和variance）更是需要额外的约280GB，加上梯度本身又是140GB——总计超过560GB的显存需求，远超任何单卡的承载能力。

传统的分布式数据并行（DDP）虽然可以解决数据分片的问题，但每张卡上仍然需要维护完整的模型副本，显存瓶颈并未突破。正是在这种背景下，微软DeepSpeed团队提出的ZeRO（Zero Redundancy Optimizer）优化技术应运而生，它通过消除数据并行中的冗余内存占用，让分布式训练能够以接近线性的扩展效率在成百上千张GPU上运行。

本文将从ZeRO的底层原理出发，深入剖析其三个阶段的内存优化机制，并结合实际的DeepSpeed配置代码，帮助你真正掌握这一大模型训练的核心技术。

深度学习分布式训练架构

ZeRO的核心思想：消除冗余，分而治之

在传统的分布式数据并行训练中，假设我们有N张GPU，每张GPU都保存了完整的模型参数、梯度、优化器状态。这意味着整个集群中，每份数据都被冗余存储了N次。ZeRO的核心理念就是：既然数据已经被分片了，为什么状态还要被冗余存储？

ZeRO通过三种不同的优化级别（Stage 1、Stage 2、Stage 3），逐步消除这些冗余，将内存占用从O(N)降为O(1/N)。

ZeRO Stage 1：优化器状态分片（Optimizer State Partitioning）

Stage 1是最基础的优化级别。它只对优化器状态（Adam的momentum和variance）进行分片，模型参数和梯度仍然每张卡各持一份。在N张卡上，每张卡只负责维护约1/N的优化器状态。

实际效果如何？以Adam优化器为例，每个参数需要存储两份优化器状态（momentum和variance），每份的大小与参数本身相同。这意味着优化器状态占用的内存大约是模型参数的两倍。Stage 1将这部分内存缩减到原来的1/N，在64卡场景下相当于将优化器状态内存从2×模型大小降低到2×模型大小/64——几乎可以忽略不计。


1
2
3
4
5
6
7
8
9
10
11
# 训练过程中每个参数的内存占用变化（FP16参数）

# 参数: 2 bytes

# 梯度: 2 bytes

# Adam states: 4 bytes (momentum fp32) + 4 bytes (variance fp32) = 8 bytes

# 总计: 12 bytes per parameter



# Stage 1 后 (N=8):

# 参数: 2 bytes (不变)

# 梯度: 2 bytes (不变)

# Adam states: 8/N = 1 byte

# 总计: 5 bytes per parameter

ZeRO Stage 2：梯度分片（Gradient Partitioning）

Stage 2在Stage 1的基础上，进一步对梯度进行分片。在反向传播过程中，每个GPU只保存自己负责的那部分梯度，其他部分的梯度在计算完成后即被丢弃。

这里有一个关键的技术细节：梯度AllReduce。在DDP中，反向传播结束后需要执行AllReduce操作来同步梯度。ZeRO Stage 2将这个过程优化为Reduce-Scatter + AllGather两个步骤。Reduce-Scatter将梯度分散归约到各卡，每卡只保留自己负责的梯度切片；AllGather则在优化器更新后，将更新后的参数广播给所有卡。


1
2
3
4
5
6
7
8
9
10
11
12
13
14
# DeepSpeed ZeRO Stage 2 配置示例（ds_config.json）

{

  "zero_optimization": {

    "stage": 2,

    "allgather_partitions": true,

    "allgather_bucket_size": 5e8,

    "overlap_comm": true,

    "reduce_scatter": true,

    "reduce_bucket_size": 5e8,

    "contiguous_gradients": true

  },

  "gradient_accumulation_steps": 8,

  "train_batch_size": 64

}

参数说明：

allgather_bucket_size：AllGather操作的缓冲区大小（字节），越大通信效率越高但内存占用也越大，推荐5e8（约500MB）
overlap_comm：是否将通信与计算重叠，开启后可以显著提升吞吐量
reduce_bucket_size：Reduce-Scatter的缓冲区大小，通常与allgather_bucket_size保持一致
contiguous_gradients：将梯度存储为连续内存块，提升通信效率

ZeRO Stage 3：参数分片（Parameter Partitioning）——训练一切模型的关键

Stage 3是ZeRO最强大的优化级别，也是训练百亿级以上模型的必备技术。在Stage 3中，不仅优化器状态和梯度被分片，模型参数也被分片到各个GPU上。这意味着在任何时刻，每张GPU只持有完整模型参数的1/N。

当某个GPU需要访问不属于它的参数时（例如在前向传播中需要计算某个层的输出），它通过AllGather从其他GPU获取所需的参数。计算完成后，这些参数副本立即被丢弃，以释放内存。

这种”按需获取”的模式带来了巨大的内存节省，但同时也引入了额外的通信开销。DeepSpeed通过以下技术来缓解这个问题：

通信与计算重叠

DeepSpeed会在前向传播中预取下一层需要的参数，在反向传播中预取下一层需要的梯度，从而实现通信与计算的流水线并行。这使得通信开销几乎被完全隐藏。


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
# DeepSpeed ZeRO Stage 3 完整配置示例

{

  "zero_optimization": {

    "stage": 3,

    "stage3_max_live_parameters": 1e9,

    "stage3_max_reuse_distance": 1e9,

    "stage3_prefetch_bucket_size": 5e7,

    "stage3_param_persistence_threshold": 1e6,

    "stage3_gather_16bit_weights_on_model_save": true,

    "sub_group_size": 1e9,

    "reduce_bucket_size": "auto",

    "overlap_comm": true

  },

  "fp16": {

    "enabled": true,

    "auto_cast": true,

    "loss_scale": 0,

    "initial_scale_power": 16

  },

  "optimizer": {

    "type": "AdamW",

    "params": {

      "lr": 3e-5,

      "betas": [0.9, 0.999],

      "eps": 1e-8,

      "weight_decay": 0.01

    }

  },

  "scheduler": {

    "type": "WarmupLR",

    "params": {

      "warmup_min_lr": 0,

      "warmup_max_lr": 3e-5,

      "warmup_num_steps": 500

    }

  }

}

参数预取策略

DeepSpeed使用了一种基于参数重用距离的预取策略。它维护一个参数使用历史记录，预测哪些参数将在不久的将来被访问，并提前发起AllGather请求。这类似于CPU的预取器，但面向的是GPU集群中的分布式参数。

参数	说明	推荐值
stage3_max_live_parameters	GPU上同时保持的最大活跃参数数量	1e9（约2GB FP16）
stage3_prefetch_bucket_size	预取缓冲区大小	5e7（约50MB）
stage3_param_persistence_threshold	小于此阈值的参数保持常驻（不释放）	1e6
sub_group_size	子组大小，用于控制通信粒度	1e9

实战：使用DeepSpeed ZeRO-3训练Llama风格模型

下面我们通过一个完整的实战案例，演示如何使用DeepSpeed ZeRO-3训练一个1.3B参数的模型。这个例子使用HuggingFace Transformers + DeepSpeed插件，是目前大模型训练最主流的方案。

环境准备


1
2
3
4
5
6
7
# 安装依赖

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

pip install deepspeed transformers datasets accelerate wandb

pip install ninja  # 加速CUDA算子编译



# 验证DeepSpeed安装

ds_report

训练脚本


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
# train_deepspeed.py

import os

import torch

import deepspeed

from transformers import (

    AutoConfig,

    AutoModelForCausalLM,

    AutoTokenizer,

    Trainer,

    TrainingArguments,

    DataCollatorForLanguageModeling,

)

from datasets import load_dataset



def main():

    # 模型配置：1.3B参数

    model_name = "facebook/opt-1.3b"

    

    # 加载分词器

    tokenizer = AutoTokenizer.from_pretrained(model_name)

    tokenizer.pad_token = tokenizer.eos_token

    

    # 加载数据集

    dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train")

    

    def tokenize_function(examples):

        return tokenizer(examples["text"], truncation=True, max_length=1024)

    

    tokenized_dataset = dataset.map(

        tokenize_function, 

        batched=True, 

        remove_columns=["text"]

    )

    

    # 训练参数

    training_args = TrainingArguments(

        output_dir="./output",

        per_device_train_batch_size=4,

        gradient_accumulation_steps=8,

        num_train_epochs=3,

        logging_steps=10,

        save_steps=500,

        save_total_limit=2,

        fp16=True,

        deepspeed="ds_config.json",  # 指向上面的ZeRO-3配置

        report_to="wandb",

        dataloader_num_workers=4,

    )

    

    # 数据整理器

    data_collator = DataCollatorForLanguageModeling(

        tokenizer=tokenizer, mlm=False

    )

    

    # 初始化模型

    model = AutoModelForCausalLM.from_pretrained(model_name)

    

    # Trainer自动加载DeepSpeed

    trainer = Trainer(

        model=model,

        args=training_args,

        train_dataset=tokenized_dataset,

        data_collator=data_collator,

        tokenizer=tokenizer,

    )

    

    trainer.train()



if __name__ == "__main__":

    main()

启动训练


1
2
3
4
5
6
7
8
# 单机多卡训练（8×A100 80GB）

deepspeed --num_gpus=8 train_deepspeed.py



# 多机训练（2节点×8卡）

deepspeed --num_gpus=8 --num_nodes=2 \

    --master_addr=192.168.1.100 \

    --master_port=29500 \

    train_deepspeed.py

ZeRO的内存节省效果量化分析

为了直观理解ZeRO各阶段的内存节省效果，我们以训练一个70B模型在64张A100 80GB上为例：

优化级别	参数量	每卡参数量	每卡内存占用	可训练最大模型
DDP（无ZeRO）	70B	70B	~840GB（溢出）	~6B
ZeRO-1	70B	70B	~770GB（溢出）	~7B
ZeRO-2	70B	70B	~700GB（溢出）	~8B
ZeRO-3	70B	70B/64≈1.1B	~13GB	~300B+

可以看到，ZeRO-3的优化效果是革命性的——它让原本需要840GB显存的70B模型训练，压缩到每卡仅需约13GB（加上激活值等额外开销，实际约需40-50GB），使得在8卡A100 80GB上训练70B模型成为可能。

ZeRO-3的高级优化技巧

在将ZeRO-3应用于实际生产环境时，以下几个优化方向值得重点关注：

1. 混合精度训练的协调

ZeRO-3与混合精度训练（FP16/BF16）有天然的协同效应。DeepSpeed的FP16优化器在ZeRO-3下会将FP32的主权重与FP16的分片参数分开管理，进一步减少内存占用。使用BF16（Brain Float 16）训练可以避免FP16的精度溢出问题，尤其适合训练大模型。


1
2
3
4
5
6
7
8
9
{

  "bf16": {

    "enabled": true

  },

  "zero_optimization": {

    "stage": 3,

    "stage3_gather_16bit_weights_on_model_save": true

  }

}

2. 卸载（Offload）到CPU/NVMe

当GPU显存仍不足以容纳模型时，ZeRO-3可以将参数卸载到CPU内存甚至NVMe固态硬盘上。这通过ZeRO-Infinity技术实现，它允许训练几乎任意大小的模型，代价是引入CPU-GPU通信延迟。


1
2
3
4
5
6
7
8
9
10
11
12
13
14
{

  "zero_optimization": {

    "stage": 3,

    "offload_param": {

      "device": "cpu",

      "pin_memory": true

    },

    "offload_optimizer": {

      "device": "cpu",

      "pin_memory": true,

      "ratio": 0.5

    }

  }

}

当CPU内存也被耗尽时，可以进一步卸载到NVMe：


1
2
3
4
5
6
7
8
9
10
11
12
{

  "zero_optimization": {

    "stage": 3,

    "offload_param": {

      "device": "nvme",

      "nvme_path": "/mnt/nvme/zero_offload",

      "pin_memory": true,

      "buffer_count": 4,

      "buffer_size": 1e8

    }

  }

}

3. 梯度检查点（Gradient Checkpointing）

Gradient Checkpointing是一种以计算换内存的技术，在前向传播中不保存中间激活值，而是在反向传播时重新计算。与ZeRO-3结合使用时，可以进一步降低激活值的内存占用。


1
2
3
4
from transformers import AutoModelForCausalLM



model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")

model.gradient_checkpointing_enable()

ZeRO-3与模型并行、流水线并行的对比

ZeRO-3并不是唯一的分布式训练范式。为了帮助读者理解不同技术的适用场景，这里做一个对比：

技术	内存减少	通信开销	适合场景
ZeRO-3	O(1/N) 参数 + 梯度 + 优化器	AllGather每层一次	Transformer类模型，同质集群
张量并行（TP）	O(1/N) 参数	AllReduce每层多次	单节点内大模型，高带宽
流水线并行（PP）	O(1/N) 参数	P2P通信，低频	跨节点，低带宽场景
FSDP（PyTorch原生）	O(1/N) 参数 + 梯度 + 优化器	与ZeRO-3类似	PyTorch生态，灵活配置

在实际生产中，大模型训练通常采用”3D并行”策略：张量并行（TP）在单节点内切分计算，流水线并行（PP）在节点间切分层，ZeRO-3在数据并行维度进一步消除冗余。DeepSpeed通过这些策略的组合，成功训练了Megatron-Turing NLG 530B等千亿级模型。

常见问题与调试技巧

通信瓶颈

如果观察到GPU利用率不高（nvidia-smi中GPU-Util低于80%），可能是通信瓶颈导致。可尝试：

增大batch size以减少通信频率
开启overlap_comm让通信与计算重叠
使用InfiniBand而非以太网（RDMA vs TCP）
调整allgather_bucket_size和reduce_bucket_size

OOM（显存溢出）

当遇到OOM时，按以下顺序排查：

降低per_device_train_batch_size
启用gradient_checkpointing
开启ZeRO的offload参数
使用activation offloading（ZeRO-3的扩展功能）
检查是否有不必要的张量留在计算图中

模型保存与加载

ZeRO-3下保存模型需要特殊处理，因为参数是分散在各卡上的：


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# 保存模型（需要gather全量权重）

from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint



# 方法1：在训练脚本中保存

trainer.save_model("./final_model")



# 方法2：从checkpoint恢复完整权重

state_dict = get_fp32_state_dict_from_zero_checkpoint(

    checkpoint_dir="./output/checkpoint-1000"

)

model.load_state_dict(state_dict)



# 推理时不需要DeepSpeed

model = AutoModelForCausalLM.from_pretrained("./final_model")

tokenizer = AutoTokenizer.from_pretrained("./final_model")

inputs = tokenizer("人工智能的未来是", return_tensors="pt")

outputs = model.generate(**inputs, max_length=100)

总结

ZeRO优化技术通过消除数据并行中的冗余内存占用，使得大规模分布式训练变得可行且高效。从Stage 1只优化优化器状态，到Stage 3的参数全分片，每个阶段都对应着不同的内存-通信权衡。在实际生产中，ZeRO-3已成为训练百亿级以上模型的标配技术。

特别值得注意的是，DeepSpeed的ZeRO与PyTorch原生的FSDP（Fully Sharded Data Parallel）在理念上非常相似，FSDP在PyTorch 2.0+中已被广泛采用。两者都实现了”参数分片+按需AllGather”的核心思想。如果你使用纯PyTorch生态，FSDP是一个很好的替代方案；如果追求更丰富的功能（如CPU/NVMe offload、混合精度控制的细粒度），DeepSpeed ZeRO-3仍然是首选。

最后，建议读者在实际项目中先在单机多卡环境下尝试ZeRO-2或ZeRO-3，逐步熟悉其配置参数和调优方法，再扩展到多机场景。训练大模型不仅需要算力，更需要对其底层优化机制有深刻的理解——希望本文能为你在这条路上提供一份实用的参考指南。