DeepSpeed项目中的ZeRO-3技术深度解析-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00103/article/details/148325618

DeepSpeed项目中的ZeRO-3技术深度解析

DeepSpeed DeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective. 项目地址: https://gitcode.com/gh_mirrors/de/DeepSpeed

什么是ZeRO优化器？

ZeRO（Zero Redundancy Optimizer）是一种革命性的内存优化技术，它通过消除数据并行过程中的内存冗余，显著提升了大规模模型训练的效率和可行性。与传统数据并行方式不同，ZeRO不再在各个GPU上复制完整的模型状态，而是将模型状态（包括优化器状态、梯度和参数）进行智能分区。

ZeRO的三个阶段

ZeRO Stage 1：优化器状态分区

仅对优化器状态进行分区
例如Adam优化器中的32位权重、一阶和二阶矩估计
每个GPU只存储和更新自己负责的那部分优化器状态

ZeRO Stage 2：梯度分区

在Stage 1基础上增加梯度分区
16位梯度也被分区存储
每个GPU只保留与自己负责的优化器状态对应的梯度

ZeRO Stage 3：参数分区

最完整的ZeRO实现
16位模型参数也被分区存储
在前向和后向传播过程中自动收集和分区参数
内存节省效果最显著

ZeRO-Infinity：突破内存限制

ZeRO-3还包含了Infinity卸载引擎，形成了ZeRO-Infinity技术，它能够：

将模型状态卸载到CPU内存
进一步卸载到NVMe存储设备
实现前所未有的内存节省
支持训练超大规模模型

配置ZeRO-3

启用ZeRO-3非常简单，只需在DeepSpeed配置文件中进行相应设置。以下是几个典型配置示例：

基础ZeRO-3配置

{
    "zero_optimization": {
        "stage": 3
    },
    "fp16": {
        "enabled": true
    }
}

优化器状态卸载到CPU

{
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu"
        }
    }
}

参数和优化器状态都卸载到CPU

{
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu"
        },
        "offload_param": {
            "device": "cpu"
        }
    }
}

卸载到NVMe存储

{
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "nvme",
            "nvme_path": "/nvme_data"
        },
        "offload_param": {
            "device": "nvme",
            "nvme_path": "/nvme_data"
        }
    }
}

构建超大规模模型

ZeRO-3使得构建参数规模远超单个节点内存容量的模型成为可能：

with deepspeed.zero.Init():
    model = MyMassiveModel()

这种上下文管理器会智能地初始化模型参数，确保它们被正确分区。

参数协调机制

当需要在模块的forward方法之外访问参数时，可以使用以下方法：

参数收集

with deepspeed.zero.GatheredParameters([param1, param2]):
    # 在这里可以访问完整的参数
    do_something(param1, param2)

注册外部参数

deepspeed.zero.register_external_parameter(module, param)

内存中心分片技术

ZeRO-Infinity包含了一项称为内存中心分片的创新技术：

将大型算子分解为可顺序执行的小分片
每个分片的参数和梯度可以单独获取和释放
显著降低工作内存需求
支持任意大小的算子，无需模型并行重构

from deepspeed.zero import TiledLinear

# 替换普通线性层
linear_layer = TiledLinear(in_features, out_features, tile_dim=1024)

调试技巧

ZeRO训练中，由于参数、梯度和优化器状态都被分区，调试可能比较复杂。DeepSpeed提供了一些实用工具来访问这些状态：

from deepspeed.utils import (safe_get_full_fp32_param, 
                           safe_get_full_grad,
                           safe_get_full_optimizer_state)

# 获取完整参数
full_param = safe_get_full_fp32_param(lp)

# 获取完整梯度
full_grad = safe_get_full_grad(lp)

# 获取完整优化器状态
optim_state = safe_get_full_optimizer_state(lp, "exp_avg")

修改分区状态

DeepSpeed也提供了修改分区参数的接口：

from deepspeed.utils import (safe_set_full_fp32_param,
                           safe_set_full_optimizer_state)

# 设置完整参数
safe_set_full_fp32_param(lp, new_tensor)

# 设置完整优化器状态
safe_set_full_optimizer_state(lp, new_state, "exp_avg")