- 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架构是什么样的，如何实现？
自动对参数进行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方法。这种替换方式就像下面这个例子：
pythonclass 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个部分：
-
初始化，module到Continer的创建流程；
- policy定义的位置； 入口：hybrid_engine.py populate_all_inference_policies policy定义文件：deepspeed/module_inject/replace_policy.py policy和container对应关系：deepspeed/module_inject/utils.py
- Container创建过程；
- Create_module 创建推理图；新的forward函数
- set_params_wo_copy container将模型变量赋值给计算图；只给引用 不复制
- forward的替换发生在eval方法里；
- 计算图的构建；
- 保证和原生pytorch的计算过程保持一致；
- 尽量使用ds提前定义的cuda算子；
-
generate的过程；
2. 如何适配新的模型
2.1 关键点
以chatglm2和BlueLM为例，几个关键点：
-
如何定义一个新的模型； 参考HE的架构，主要新增policy、container，以及使用对应的ops重构计算图；
- Policy：
- 定义原模型中关键的参数变量；
- 定义原模型中的参数引用；
- 定义要替换的模块，一般为transformer层级模块对应的类；
- Container；
- 保存计算图中需要的参数tensor引用；
- 生成推理用的module，用module.forward替换掉原模块中的forward；
- 重构推理计算图；
- 检查算子是否可以复用；
- 检查结果是否正确；
- Policy：
-
TP相关的代码；
改动代码之前，请阅读arxiv.org/pdf/1909.08...论文，了解Tensor Parallelism的大模型分布式训练方式。在HE中只关注前向推理的流程，适配时最主要的是确定好参数切分的方式。
- 参数复制，拆分；
- 推理时的reduce操作；
-
计算流程实现；
-
*Inference相关的代码；
chatglm2 66b 8卡推理；
2.2 适配开发流程
HE本质上就是用zero3参数收集起来，重新生成一个计算图（container.module.forward函数）复现并替换原pytorch模型的推理结果；
-
找到要替换的类，policy对应的类； 一般在要适配的模型定义文件里面；
-
找到container.module.forward调用入口； 在替换模型定义文件里，policy对应的类的forward函数的调用入口，注意入口的参数传递和container对应的forward声明对齐；
-
检查点：
- 参数引用是否正确；
- 计算流程是否一致；
-
单步调试； 每一步的计算和pytorch代码是否一致；
- TP代码开发的区别：
- container需要单独保存参数； 原始参数的切片；
- 在生成阶段，需要先准备tp参数；更重要的是，在生成结束之后要及时释放tp的参数； 在container的有相应的方法需要实现，主要有参数准备、参数释放；
- 计算阶段需要适配tp流程； 显然的变化是，计算用的weight的大小会改变；
- 注意：generate阶段需要对输入进行拼接；
2.3 单元测试方法
-
如何进行多机程序的debug；
主要难点是多机程序的调试，这里提供2种方式；
-
可以使用pdb在单节点上进行调试代码；代码如下：
python## 在需要阻塞的位置插入下面的代码 if dist.get_rank() == 0: # 判断是否为0的进程 breakpoint() # 断点 dist.barrier() # 当0进程阻塞时，其他节点进行等待 -
（理论可行，未做尝试）用单进程模拟多进程；
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； 使用提前开辟好内存的方式
- BlueLMAttention细节：
==修复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的时候超出限制；
- 解决方法如下：
- 修改自定义算子申请显存的逻辑； 代码在csrc/transformer/inference/includes/inference_context.h#L98 GenWorkSpace函数里
调大ds_bluelm.py里面申请显存的参数；代码在deepspeed/model_implementations/transformers/ds_bluelm.py- BlueLMAttention里forward计算attention的部分改成pytorch实现； pytorch动态申请显存，只要显存够就不存在溢出的风险；
-
第2种方法，通过调大长度参数，仍然没有解决；
-
第3种方法，使用pytorch重新定义attention计算过程；
-
kv-cache适配时，遇到present_key_value向量无理由自发改变的问题；
通过改变present_key_value内存地址解决；==原因未知，如果有，就是对cuda敬畏之心不够！==
-