李白:你的模型权重很不错,可惜被我没收了
撰文|CPFLAME
大噶好,年更楼主今天想推的是,主打分布式训练的模型库_李白(LiBai)。
https://github.com/Oneflow-Inc/libai
对于目前市面上的模型库来说,选择实在是太多了,换了一批又一批,眼睛都挑花了,为什么要用 LiBai?(如果你觉得 LiBai 万一某天能用到,或者这篇文章读下来感觉比较开心,可以去 GitHub 上点赞,如果能三连就更好了。众所周知,GitHub 点赞其实是个收藏夹功能)。
按照现在的趋势来说,模型越来越大了,大到一张 GPU 甚至装不下完整的模型,必须得上分布式并行技术,但是分布式代码在很多框架下都是高度定制化的,对于新手来说根本读不懂,也不知道应该怎么使用,导致大家上手非常的困难,让自己珍贵的发际线显得更加珍贵。
针对大模型上述存在的痛点,导致我们必须上分布式(数据并行、模型并行、流水并行)才能跑起来一个大模型。
那么,LiBai 有哪些特点呢?你坐好,我要发功了。
需要详细分章介绍的优势(看上去还不错,用户也可以听得懂,也知道要干什么):
简单易用的分布式代码,单机代码和分布式代码基本一致
可以无缝使用 PyTorch、HuggingFace 的 model 权重,并且还可以在 LiBai 下进行多机多卡的分布式推理
开箱即用,所有的分布式并行配置(Grad Acc,AMP,Checkpointing,ZeRO,Auto Parallel)技术都只需要在 config 里面一键设置就可以生效,不需要在算法代码 model.py 中额外添加
支持模型一键转换 ONNX
我搁这儿就要介绍完的优势(看上去大家也有,很虚的帽子话),为了不让大家觉得过于虚,在介绍的同时也会插入相关的例子。
1. 具有高度灵活性和高效率,同时支持动态图 eager 模式和静态图 graph 模式,支持一键切换,在方便 debug 和高效性之间反复横跳。
2. 对于分布式并行的支持比较全面,大家可以在里面尽情地组合各种分布式并行的组件。
3. LiBai 下面有内置的 layers 直接使用,避免重复造轮子,比如用 LiBai 下面的 Linear 层就可以快速地构建一个 2D 并行(数据并行+模型并行)的 MLP。
4. 采用 LazyCall(借鉴自 detectron2)的配置系统,基于 Python 语法构建相比于 argparse 和 yacs-based 方式更灵活,而且每次训练都会序列化 yaml 文件,用户可以一键读取 yaml 文件来复现“上古时期”的实验结果。
5. 具有丰富的 Projects 实现。由于 LiBai 的分布式并行设计与算法逻辑进行了解耦,使得在 Projects 下面的算法都可以享受到 LiBai 下面的分布式并行技术,而且随着分布式并行技术的更新,Projects 下面的算法代码不需要任何更新就可以享受到更新后的成果。
6. 和业界翘楚 Megatron 比起来,具有不弱于它的吞吐,甚至稍占优势,完整的对比实验在 LiBai tutorial(https://libai.readthedocs.io/en/latest/tutorials/get_started/Benchmark.html)和《大模型训练难于上青天?效率超群、易用的李白模型库来了》,这里给一个 GPT2 的 3D 并行数据简单感受一下。
可能有人会问,怎么都和 Megatron 去比,你们各个同行之间有对比数据吗?主要原因有二:1) 只要吕布不说话反驳,那么我邢道荣就有不下于吕布的勇武,人均小吕布,这很合理;2) 大家都是国产框架,中国人不卷中国人,咱们薅着一个外国人可劲儿的打。
下面分章详细说说上述优点。
开源自助
LiBai 最基础的一个功能: 开源自助。也就是除了 LiBai 训练出来的模型以外,我们还可以加载 PyTorch 以及 HuggingFace 上面的模型进行分布式推理。
由于 LiBai 的底层是基于 OneFlow 来实现的,而 OneFlow 的算子绝大部分都已经和 PyTorch 进行了对齐,这能发挥出什么优势?
使用前还要看看预备知识。一个完整的模型由两个部分构成:模型结构,换种说法就是 model.py 模型权重,再换种说法就是 model_best.pth。
假设我们在框架 A 下面,有 modelA.py 和 model_best_A.pth,我们想在框架 B 上面跑起来这个框架 A 下面的模型,应该怎么做呢?
在框架 B 下面,用框架 B 的算子搭建出一个 modelB.py,该 modelB 的参数名字可以和 modelA 的不一致,但是前向推理的逻辑运算最好一致,然后去加载 model_best_A.pth 得到 model_A_state_dict(),把 model_A_state_dict()里面的参数格式全部转换成框架 B 下面支持的格式,其中可以运用中间格式进行转换。
举个例子,比如 torch.tensor()->np.numpy()(中间格式)->oneflow.tensor()之前提到了 modelB 中的参数名字可以和 modelA 中的不一致,如果不一致,那么需要把 model_A_state_dict()中的 key 值改一下和 modelB 的一致。
做完了以后,直接加载我们转换好的参数
modelB.load_state_dict(model_A_state_dict),就可以在框架 B 下面进行推理。为了保证模型转换好以后的准确性,可以喂给 modelA 以及 modelB 相同的输入,检查一下是否能得到相同输出。
这个预备知识不仅限于 LiBai,在任何模型复现或者模型迁移上面都适用。
有了预备知识以后怎么使用 PyTorch 或者 HuggingFace 下面的模型?简单来说分为以下几步:
把 torch 的算子替换为 oneflow: 把 torch_model.py 下面的 torch 全部替换为 oneflow,得到 oneflow_model.py.把 oneflow_model.py 中的 layer 尽可能地替换成 LiBai 中支持的 layer,只替换你想要的部分也可以(比如只替换 Linear 层),LiBai 会自动把没有替换的 layer 转换成分布式并行所需要的格式。
这一步是支持分布式推理的关键继承 LiBai 提供好的分布式推理的基类 basic.py,重载转换权重的函数,按照 PyTorch 那样写好预处理和后处理,就可以进行分布式推理了。
下面链接里面有极其详细的步骤解答,看在作者不仅授人以鱼,还授人以渔,还做了程序员最讨厌的文档活儿,可以顺便给 LiBai 点个 star 收藏,而且保不齐以后万一有个什么复现的任务,这里面的知识点也用得上,至少可以用别人 release 出来的预训练权重来验证自己复现的 model.py 是否正确。这叫什么?这叫 call back!
LiBai 分布式推理介绍:
http//:github.com/Oneflow-Inc/libai/discussions/386
以 MT5 应用 HuggingFace model 为例子,我们在 2 机 4 卡下面进行模型并行 2X 流水并行 2 的分布式推理,跑起来的代码风格如下:
那么多机多卡的分布式推理脚本
在 node0 上输入指令:
在 node1 上输入指令:
细心的朋友已经发现了,LiBai 下面可以通过设置 pipeline_stage_id, 来让用户自己设置每个 stage 上 group 的层数是多少,方便在某些极端情况下(比如你的机器 0 很强,但是机器 1 很拉胯,或者你的 encoder 计算量巨大,但是 decoder 计算量较小)手动实现负载均衡。
大模型训练
众所周知,大家都喜欢做点"出格"的事情,比如在上班的时候摸鱼,在 VScode 上面炒股......
那么 LiBai 呢?你甚至可以拿它来训练模型!
在 Projects(https://github.com/Oneflow-Inc/libai/tree/main/projects)下支持的模型:
下面来谈谈模型之外,LiBai 有什么不一样的地方,换句话说,也就是核心竞争力在哪里?
分布式配置和算法逻辑解耦
LiBai 进行了模块化的设计,使得分布式的配置和算法逻辑解耦,这意味着什么?
这意味着用户只需要把大部分的注意力专注到算法逻辑上面,而不用再苦恼怎么插入各种并行的代码了。
简单来说,下面这些模块都可以在 config.py 中进行一键配置。
单机和分布式代码几乎一致
下面给一个简单的 2D 并行(数据并行+模型并行)的 MLP 例子, 比如你的 Linear 层在 16384 这个维度上面比较大, 需要把它切分在不同的卡上才能装下, 那么在 LiBai 下面只需要如下所示就可以完成了,几乎跟单机代码没有区别。
支持一键转换 ONNX
本人对一键转 ONNX 的执念可谓是相当之深了。同样以 MT5 为例子,LiBai 支持了一键转换 ONNX 的功能,点击以下链接就可以体验:
https://github.com/Oneflow-Inc/libai/tree/main/libai/onnx_export
更详细的说明和教程会在 LiBai 中持续发布。如果这篇文章对你有启发,请不要吝惜手中的收藏按钮,欢迎去 GitHub 上 Star、Fork、Watch 三连,持续跟进最新进展。
GitHub 地址:https://github.com/Oneflow-Inc/libai
也欢迎你扫码进技术交流群探讨。
引用
1. https://github.com/facebookresearch/detectron2
2. https://github.com/Oneflow-Inc/oneflow
3. https://github.com/Oneflow-Inc/oneflow_convert
4. https://github.com/NVIDIA/Megatron-LM
其他人都在看
欢迎下载体验 OneFlow v0.8.0 最新版本:https://github.com/Oneflow-Inc/oneflow/
评论