数据并行与ZeRO

数据并行回顾

数据并行（Data Parallelism，DP）是最基础的分布式训练范式：

模型复制：每张GPU持有完整模型副本
数据分片：不同GPU处理不同数据batch
梯度同步：每轮迭代后对各GPU的梯度做AllReduce同步
参数更新：所有GPU使用相同的更新后参数继续下一轮

GPU 0: [模型副本] + [数据分片0] → [本地梯度] → AllReduce → [更新后模型]
GPU 1: [模型副本] + [数据分片1] → [本地梯度] → AllReduce → [更新后模型]
...

朴素数据并行的局限

显存浪费：每张GPU保存完整模型、激活值、优化器状态，显存占用极大
通信瓶颈：梯度AllReduce通信量 = 模型参数量，对大模型不现实
扩展性差：模型参数量超过单卡显存时无法运行

ZeRO（Zero Redundancy Optimizer）

ZeRO是微软DeepSpeed团队提出的显存优化技术，通过分片（sharding）消除数据并行中的显存冗余。

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

原理：将Adam优化器的状态（动量、方差）在数据并行进程间分片，每个GPU只保存 1/N 的优化器状态。

优化器状态总量: 2 * 参数总量（动量 + 方差）
Stage 1后: 每GPU只保存 2 * 参数总量 / N

通信开销：

通信量：与标准数据并行相同（梯度AllReduce）
额外通信：参数分片 scatter/gather（可选）

收益：约4倍显存降低（假设N=8，则显存降至1/4）

Stage 2：梯度分片（Gradient Sharding）

原理：在Stage 1基础上，进一步将梯度分片存储。

标准DP: 每GPU保存完整梯度 → AllReduce → 每GPU有完整梯度 → 参数更新
Stage 2: 每GPU保存 1/N 梯度 → 需要时 AllReduce（分片通信）

实现机制：

反向传播时本地计算梯度
执行Reduce-Scatter，每个GPU获得其他GPU数据的部分和
本地更新对应参数分片
下次前向传播前执行All-Gather重建完整梯度（可选）

通信模式：

通信量：与Stage 1相同，但通信模式不同（从AllReduce变为Reduce-Scatter + AllGather）
可与通信计算重叠进一步优化

收益：约8倍显存降低（结合Stage 1）

Stage 3：参数分片（Parameter Sharding）

原理：将模型参数也分片存储，每个GPU只保存 1/N 的参数。

完整分片流程:
前向/反向: 动态 All-Gather 获取所需参数分片 → 计算 → 丢弃非必要分片

实现机制：

前向传播时：参数分片本地计算，需要邻 GPU 参数时通过All-Gather获取
反向传播时：梯度分片本地更新，触发Reduce-Scatter同步
每轮迭代结束后参数分片自动同步

通信开销：

每层前向/反向都需要 All-Gather（获取参数）和 Reduce-Scatter（同步梯度）
通信量显著增加，但换取极大显存节省

收益：理论上显存降低与GPU数成线性关系

ZeRO 通信量分析

Stage	通信量（相比DP）	显存收益
Stage 1	1x (相同)	4x
Stage 2	1x (相同)	8x
Stage 3	~2x (AllGather+ReduceScatter)	O(N)

ZeRO-DP 与标准DP对比

标准DP:
每GPU内存占用 = 参数 + 梯度 + 优化器状态 + 激活值
            ≈ 模型大小 × (1 + 2 + 4) ≈ 7× 模型大小（16位存储）

ZeRO-3:
每GPU内存占用 ≈ (参数 + 梯度 + 优化器状态) / N + 激活值
            ≈ 模型大小 / N + 激活值

DeepSpeed实现

配置文件示例

{
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": true
        },
        "overlap_comm": true,
        "contiguous_gradients": true,
        "reduce_bucket_size": 1e8,
        "stage3_prefetch_bucket_size": 1e9,
        "stage3_param_persistence_threshold": 1e5
    },
    "bf16": {
        "enabled": true
    }
}

核心API

from deepspeed import DeepSpeedConfig, initialize
model, optimizer, _, _ = initialize(
    args=args,
    model=model,
    model_parameters=model.parameters(),
    config=deepspeed_config
)

ZeRO-3 关键技术

分片通信与计算重叠：利用CUDA stream在计算时通信
Gradient Accumulation：梯度累积避免流水线空闲
Weight Gradients Partitioning：参数梯度自动分片
Partition Communication Group：自定义通信域

ZeRO-Stage 3与PipeDream的结合

DeepSpeed支持将ZeRO与流水线并行结合：

{
    "zero_optimization": {"stage": 3},
    "pipeline": {"stages": 4}
}

每个pipeline stage内部使用ZeRO-3分片，进一步降低单卡显存。

Offload技术

DeepSpeed支持将部分数据卸载到CPU或NVMe：

ZeRO-Offload：优化器状态卸载到CPU，配合Stage 2
ZeRO-Infinity：参数和优化器状态卸载到NVMe，支持超大规模模型

性能调优

通信瓶颈：TP/PP机器与DP机器分离，避免跨域通信
通信域配置：优化 NCCL 组配置
Mixed Precision：BF16训练减少通信量
Async IO：参数卸载异步化避免阻塞

与其他并行策略的关系

策略	维度	分片对象
张量并行	层内	参数/激活切分
流水线并行	层间	模型层分Stage
数据并行	batch	梯度/参数同步
ZeRO	数据并行维度	参数/梯度/优化器状态分片

三者可组合：TP=8, PP=4, ZeRO-3 + DP = N/(8*4)

框架对比

特性	DeepSpeed ZeRO	Megatron-FSDP	PyTorch FSDP
分片方式	优化器状态+梯度+参数	参数+梯度+优化器状态	参数+梯度
Offload	CPU/NVMe	CPU	CPU
通信优化	Overlap + Partition	Async	Overlap
易用性	高（配置驱动）	中	中

Huajishe Notes

Explorer

3-数据并行与ZeRO