DeepSpeed-HybridEngine开发指南

陈道一2024-01-07 14:37
  • 2023-11-29写；
  • 2023-12-06修改：增加适配模型开发流程说明；增加bug解决记录；

Deepspeed-Chat是一个优秀且易用的PPO开源实现，实际在使用时HybridEngine开发是PPO工程相关的重要一环，本次分享的目的：

  • 了解整体Deepspeed的架构，和代码逻辑；
  • 清楚如何在Deepspeed上进行HE适配相关开发；
    • 主要是对新的模型结构的适配；

先回答为什么需要做适配，因为HE(hybrid engine)本身解决的问题是将训练中的zero3模型，转换成更高效的对设备通讯压力不大的推理模式，可以是不带tp的全参数推理，也可以是带tp的推理。所以不管是哪种形式的推理，都需要重构推理图，并且处理好这种模式转换间设计的大量引用，也是适配需要做的全部事情。虽然我相信这个过程一定可以改成完全自动的形式，但是目前还没有找到这种实现方式。

然后看一下HE的优势是什么，推理速度不用说，就相当于带或不带zero3的差距，通常10x左右。还有一个优势是带TP的HE方式在内存上的优势，这一点在模型比较大且显存压力较大的场景下尤为重要。下面举ppo的例子，如果训练的sft和actor模型大小是70b，reward和critic是7b，一共4个模型，考虑32张A100-80G的场景：

  • ZeRO3训练模型占用显存（每张卡）：

    Actor: (优化器12 + 参数2 + 梯度2) * 70 / 32 = 35G SFT：通常offload，显存可以记作0G reward+critic：3.9G 每张卡总共需要：38.9G

    占用的显存是够的，但是zero3推理速度在ppo生成阶段几乎不可用；

  • （HE模式）ZeRO3训练模型+TP推理占用显存（每张卡）：

    TP的size设为8

    • 训练阶段：TP的参数释放掉，和ZeRO3模式一样，38.9G；
    • 生成阶段：70 * 2 = 140G（全部推理参数），TP参数切片140 / 8 = 17.5G，所以一共需要56.4G；

    训练用ZeRO3节约显存，生成用HE提升速度。如果不用TP的HE那么140G放到一张卡上，目前还没有设备能支持。

1. 整体架构

1.1 启动流程

涉及代码都在deepspeed/launcher目录下面，一共2个文件 runner和launcher。

1.2 Zero3

  • zero3架构是什么样的，如何实现？

    github.com/microsoft/D...

    自动对参数进行all_gather，使用后自动释放；实现核心 ZeROOrderedDict

1.3 Hybrid Engine

  • 为什么需要HE？

    Zero3是用来训练的并行方式，推理的时候有很大劣势，并且不能扩展到多机的情况。

  • HE如何起到作用？实现方式？

    HE内容：1. 自定义算子；2. tensor parallelism；

    • 不带TP的方式； 适用于 7b、13b 单机多卡
    • 带TP的方式； 适用于66b、70b 更大模型的训练

    比如说模型ChatGLM2，它的代码实现里transformer块对应的实现类是 GLMBlock，HE就是实现一个新的推理过程，带或者不带TP区别就是这个推理过程是不是分布式，然后替换掉 GLMBlock的forward方法。

    这种替换方式就像下面这个例子：

    python 复制代码 
    class Bird:
    def move(self,):
        print(f'i\'m flying')

class Pig:
    def move(self,):
        print(f'i\'m walking')

bird = Bird()
pig = Pig()

pig.move = bird.move
pig.move()
''' output
>>> i'm flying
'''

    在python里，Pig想要飞的方式是借用Bird的翅膀，而不需要通过继承实现。

运作流程

Container、Inference、Module、ops之间的调用关系；

HE运作流程，有2个部分：

  1. 初始化，module到Continer的创建流程；

    1. policy定义的位置； 入口：hybrid_engine.py populate_all_inference_policies policy定义文件：deepspeed/module_inject/replace_policy.py policy和container对应关系：deepspeed/module_inject/utils.py
    2. Container创建过程；
      1. Create_module 创建推理图；新的forward函数
      2. set_params_wo_copy container将模型变量赋值给计算图；只给引用 不复制
      3. forward的替换发生在eval方法里；
    3. 计算图的构建；
      1. 保证和原生pytorch的计算过程保持一致；
      2. 尽量使用ds提前定义的cuda算子；

  2. generate的过程；

2. 如何适配新的模型

2.1 关键点

以chatglm2和BlueLM为例，几个关键点：

  1. 如何定义一个新的模型； 参考HE的架构，主要新增policy、container，以及使用对应的ops重构计算图；

    1. Policy：
      1. 定义原模型中关键的参数变量；
      2. 定义原模型中的参数引用；
      3. 定义要替换的模块，一般为transformer层级模块对应的类；
    2. Container；
      1. 保存计算图中需要的参数tensor引用；
      2. 生成推理用的module，用module.forward替换掉原模块中的forward；
    3. 重构推理计算图；
      1. 检查算子是否可以复用；
      2. 检查结果是否正确；

  2. TP相关的代码；

    改动代码之前，请阅读arxiv.org/pdf/1909.08...论文，了解Tensor Parallelism的大模型分布式训练方式。在HE中只关注前向推理的流程，适配时最主要的是确定好参数切分的方式。

    1. 参数复制，拆分；
    2. 推理时的reduce操作；

  3. 计算流程实现；

  4. *Inference相关的代码；

    chatglm2 66b 8卡推理；

2.2 适配开发流程

HE本质上就是用zero3参数收集起来，重新生成一个计算图（container.module.forward函数）复现并替换原pytorch模型的推理结果；

  1. 找到要替换的类，policy对应的类； 一般在要适配的模型定义文件里面；

  2. 找到container.module.forward调用入口； 在替换模型定义文件里，policy对应的类的forward函数的调用入口，注意入口的参数传递和container对应的forward声明对齐；

  3. 检查点：

    • 参数引用是否正确；
    • 计算流程是否一致；

  4. 单步调试； 每一步的计算和pytorch代码是否一致；

  • TP代码开发的区别：
    • container需要单独保存参数； 原始参数的切片；
    • 在生成阶段，需要先准备tp参数；更重要的是，在生成结束之后要及时释放tp的参数； 在container的有相应的方法需要实现，主要有参数准备、参数释放；
    • 计算阶段需要适配tp流程； 显然的变化是，计算用的weight的大小会改变；
    • 注意：generate阶段需要对输入进行拼接；

2.3 单元测试方法

  • 如何进行多机程序的debug；

    主要难点是多机程序的调试，这里提供2种方式；

    1. 可以使用pdb在单节点上进行调试代码；代码如下：

      python 复制代码 
      ## 在需要阻塞的位置插入下面的代码
if dist.get_rank() == 0:   # 判断是否为0的进程
    breakpoint()   # 断点
dist.barrier()  # 当0进程阻塞时，其他节点进行等待

    2. （理论可行，未做尝试）用单进程模拟多进程；

      deepspeed启动命令中提供了 --force_multi参数，使单进程可以用来模拟多进程，并在vscode中配置启动命令用deepspeed，这样可以使用vscode进行单步调试。

    pdb的了解和相关命令查看pdb --- Python 的调试器 --- Python 3.12.0 文档。

  • 调试程序如何写；

    参考deepspeed/tests/下面对应的文件，比如：test_ds_bluelm.py。

3. 开发中BUGs解决过程

INFLIGHT状态异常

耗时：2天

  • INFLIGHT状态参数出现异常，解决思路：

    • 验证inflight状态出现的时机；

      • 分别在生成训练之后打印参数的状态；

        text 复制代码 
        第1次生成...
Counter({<ZeroParamStatus.NOT_AVAILABLE: 2>: 925})
第1次生成...
第1次训练...
Counter({<ZeroParamStatus.NOT_AVAILABLE: 2>: 762, <ZeroParamStatus.AVAILABLE: 1>: 163})
第1次训练...
第2次生成...
Counter({<ZeroParamStatus.NOT_AVAILABLE: 2>: 758, <ZeroParamStatus.AVAILABLE: 1>: 163, <ZeroParamStatus.INFLIGHT: 3>: 4})
发生报错

        流程如上面的记录，在第二次训练的时候有参数没有释放；

      • 查看训练过程中，AVAILABLE状态的变量有哪些？

        • Available 变量：

          python 复制代码 
          'model.layers.x.self_attn.k_proj.lora_left_weight',
 'model.layers.x.self_attn.v_proj.lora_left_weight',
 'model.layers.x.input_layernorm.weight',
 'model.layers.x.post_attention_layernorm.weight'
40层 * 4 = 160个 
 'model.norm.weight',
 'model.embed_layer_norm.weight',
 'model.embed_layer_norm.bias'
3个

        • Inflight 变量

          python 复制代码 
          'model.layers.39.self_attn.q_proj.weight',
 'model.layers.39.self_attn.q_proj.lora_right_weight',
 'model.layers.39.self_attn.q_proj.lora_left_weight',
 'model.layers.39.self_attn.k_proj.weight'

        大概率和lora是有关系的，从这里下手找原因；

      • 修复container中lora参数的bug；

      • 修复mlp参数中lora没有起作用；

      • 修复mlp和attention中参数拷贝的bug；

        • BlueLMAttention细节：
          • self.attn_qkvw 需要等于None；
          • self.attn_qw, self.attn_kw, self.attn_vw 在推理过程中合并成qkvw； 使用提前开辟好内存的方式

      ==修复lora相关代码后，问题仍然存在。==

    • 测试glm2-6b用zero3训练的时候，参数状态变化情况；

      python 复制代码 
      生成阶段
Counter({<ZeroParamStatus.NOT_AVAILABLE: 2>: 423})

训练阶段
Counter({<ZeroParamStatus.NOT_AVAILABLE: 2>: 338, <ZeroParamStatus.AVAILABLE: 1>: 85})

训练后第一次生成阶段
Counter({<ZeroParamStatus.NOT_AVAILABLE: 2>: 337, <ZeroParamStatus.AVAILABLE: 1>: 85, <ZeroParamStatus.INFLIGHT: 3>: 1})
['transformer.encoder.layers.0.input_layernorm.weight', 'transformer.encoder.layers.0.self_attention.query_key_value.bias', 'transformer.encoder.layers.0.post_attention_layernorm.weight', 
 ... # 每一层
 'transformer.encoder.final_layernorm.weight']
['transformer.encoder.layers.27.mlp.dense_4h_to_h.weight']

      和bluelm测试的效果差不多，但是glm6b可以正常训练；

      • 检查glm6b的代码有什么区别；

        • 训练脚本； gradient checkpoint、offload不一致
        • container 参数复制；

        ==去掉gradient checkpoint之后可以正常训练；==

    • 梳理zero3过程，找到为什么会出现inflight状态，以及对应的意义；

      • AVAILABLE 和 INFLIGHT 状态转换图
      • INFLIGHT # parameters are being gathered.

显存超出bug

耗时：一天

  • 显存超出BUG解决；

    • 原因分析：显存workspace申请不够，导致在使用算子计算softmax的时候超出限制；
    • 解决方法如下：
      1. 修改自定义算子申请显存的逻辑； 代码在csrc/transformer/inference/includes/inference_context.h#L98 GenWorkSpace函数里
      2. 调大ds_bluelm.py里面申请显存的参数； 代码在deepspeed/model_implementations/transformers/ds_bluelm.py
      3. BlueLMAttention里forward计算attention的部分改成pytorch实现； pytorch动态申请显存，只要显存够就不存在溢出的风险；

  • 第2种方法，通过调大长度参数，仍然没有解决；

  • 第3种方法，使用pytorch重新定义attention计算过程；

    • kv-cache适配时，遇到present_key_value向量无理由自发改变的问题；

      通过改变present_key_value内存地址解决；==原因未知，如果有，就是对cuda敬畏之心不够！==

相关推荐
hsling松子2 小时前
使用PaddleHub智能生成，献上浓情国庆福
人工智能·算法·机器学习·语言模型·paddlepaddle
正在走向自律2 小时前
机器学习框架
人工智能·机器学习
dengqingrui1233 小时前
【树形DP】AT_dp_p Independent Set 题解
c++·学习·算法·深度优先·图论·dp
C++忠实粉丝3 小时前
前缀和(8)_矩阵区域和
数据结构·c++·线性代数·算法·矩阵
好吃番茄3 小时前
U mamba配置问题；‘KeyError: ‘file_ending‘
人工智能·机器学习
ZZZ_O^O3 小时前
二分查找算法——寻找旋转排序数组中的最小值&点名
数据结构·c++·学习·算法·二叉树
CV-King4 小时前
opencv实战项目（三十）：使用傅里叶变换进行图像边缘检测
人工智能·opencv·算法·计算机视觉
禁默4 小时前
2024年计算机视觉与艺术研讨会（CVA 2024）
人工智能·计算机视觉
代码雕刻家4 小时前
数据结构-3.9.栈在递归中的应用
c语言·数据结构·算法
雨中rain4 小时前
算法 | 位运算（哈希思想）
算法
热门推荐
01【经验分享】Ubuntu22.04安装微信(linux官方版)02安卓系列机型永久去除data分区加密 详细步骤解析032024年高教社杯数学建模国赛C题超详细解题思路分析04CTF网络安全大赛简单的web抓包题目：HEADache05组基轨迹建模 GBTM的介绍与实现（Stata 或 R）06【2024数模国赛赛题思路公开】国赛B题思路丨附可运行代码丨无偿自提07pve(Proxmox VE)安装i225v网卡驱动08苍穹外卖面试总结09【2024高教社杯全国大学生数学建模竞赛】B题 生产过程中的决策问题——解题思路 代码 论文10RAG 实践- Ollama+RagFlow 部署本地知识库