李白: 你的模型权重很不错, 但可惜下一秒就是我的了

前言

大噶好, 年更楼主今天想推的是, 主打分布式训练的模型库_李白 (https://github.com/Oneflow-Inc/libai ), 白嫖的白.

看过楼主之前文章的老观众都会知道, 这篇稿子会很骚, 会很不正经, 但是通篇下来不会有捧一踩一, 会在表情包的轰炸和骚里骚气的言语中突出李白的特点.

所以如果大家觉得LiBai万一某天能用到, 或者这篇文章读下来也非常开心, 或者文中表情包下一秒成为你的了, 可以去github上点赞,如果能三连就更好了. 众所周知, github点赞其实是个收藏夹功能, 大家点完赞相当于把放到收藏夹里面吃灰.

https://github.com/Oneflow-Inc/libai

背景

对于目前市面上的模型库来说, 选择实在是太多了, 换了一批又一批, 眼睛都挑花了, LiBai是什么吊毛, 为什么要用LiBai?

按照现在的趋势来说, 模型越来越大了, 大到一张GPU甚至装不下完整的模型, 必须得上分布式并行技术, 但是分布式的代码在很多框架下面都是高度定制化的,对于新手来说根本读不懂, 也不知道应该怎么使用. 导致大家上手非常的困难, 自己珍贵的发际线更加珍贵.

不要脸的自卖自夸

针对大模型的痛点, 由于显存的限制, 导致我们必须上分布式(数据并行, 模型并行, 流水并行)才能跑起来一个大模型

那么LiBai具有哪些特点呢, 下面大家坐好, 我要发功了.

• 需要详细分章介绍的优势(看上去还不错, 用户也可以听得懂, 也知道要干什么):

1. 简单易用的分布式代码, 单机代码和分布式代码基本一致.
2. 可以白嫖pytorch, hugging face的model权重, 并且嫖完后可以在LiBai下面进行多机多卡的分布式推理.
3. 开箱即用, 所有的分布式并行配置(grad acc, amp, checkpointing, ZERO, auto parallel)技术都只需要在config里面一键设置就可以生效, 不需要在算法代码model.py中额外添加.
4. 支持模型一键转换onnx

• 我搁这儿就要介绍完的优势(看上去大家也有, 很虚的帽子话), 为了不让大家觉得过于虚, 在介绍的同时也会插入相关的例子

1. 具有高度灵活性和高效率, 同时支持动态图eager模式和静态图graph模式, 支持一键切换, 在方便debug和高效性之间反复横跳.

2. 对于分布式并行的支持比较全面, 大家可以在里面尽情的组合各种分布式并行的组件.

3. LiBai下面有内置的layers直接使用, 避免重复造轮子, 比如我们用LiBai下面的Linear层就可以快速的构建一个2D并行(数据并行+模型并行)的MLP.

4. 采用LazyCall(借鉴自detectron2: https://github.com/facebookresearch/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 和 微信PR稿：https://mp.weixin.qq.com/s/Xrbz0PV_kciAQI61v5JmFQ ,这里给一个GPT2的3D并行数据简单感受一下.可能有人会问了, 怎么你们都和Megatron去比, 你们各个同行之间有对比数据吗, 主要原因有二: 1) 只要吕布不说话反驳,那么我邢道荣就有不下于吕布的勇武, 人均小吕布, 这很合理. 2) 大家吃的都是国产框架的皇粮, 中国人不卷中国人, 咱们薅着一个外国人可劲儿的打.

详细说说我们需要分章来介绍的优点

快乐白嫖

重点讲一下LiBai最快乐的一个功能: 白嫖. 也就是 除了LiBai训练出来的模型以外, 我们还可以加载pytorch以及hugging face上面的模型进行分布式推理

由于LiBai的底层是基于oneflow来实现的, 而oneflow的算子绝大部分都已经和pytorch进行了对齐,那么oneflow算子和pytorch的算子进行对齐以后, 能发挥出什么优势? 在这一刻我福如心至, 这他妈不是方便我进行我最爱的白嫖了吗?

白嫖前的 预备知识 :

一个完整的模型由两个部分构成:

• 模型结构, 换种说法就是model.py
• 模型权重, 换种说法就是model_best.pth

那么假设我们在框架A下面, 有modelA.py和model_best_A.pth, 我们想在框架B上面跑起来这个框架A下面的模型, 应该怎么做呢?

1. 在框架B下面, 用框架B的算子搭建出一个modelB.py, 该modelB的参数名字可以和modelA的不一致, 但是前向推理的逻辑运算最好一致
2. 然后去加载model_best_A.pth得到model_A_state_dict(), 把model_A_state_dict()里面的参数格式全部转换成框架B下面支持的格式, 其中可以运用中间格式进行转换. 举个例子, 比如torch.tensor()->np.numpy()(中间格式)->oneflow.tensor()
3. 之前提到了modelB中的参数名字可以和modelA中的不一致, 如果不一致的话, 那么我们需要把model_A_state_dict()中的key值改一下和modelB的一致.
4. 做完了以后, 直接加载我们转换好的参数modelB.load_state_dict(model_A_state_dict, ), 就可以在框架B下面进行推理了.
5. 为了保证模型转换好以后的准确性, 可以喂给modelA以及modelB相同的输入, 检查一下是否能得到相同的输出.

这个预备知识不仅限于LiBai, 在任何模型复现或者模型迁移上面都适用. 也就是说掌握了这一套知识. 你将成为白嫖界的王者.我的儿子，你出生的那一天，整个洛丹伦的森林都在低语着一个名字——二傻子

那么有了预备知识以后我们怎么白嫖pytorch或者hugging face下面的模型? 简单来说分为以下几步:

1. 把torch的算子替换为oneflow: 把torch_model.py下面的torch全部替换为oneflow, 得到oneflow_model.py.
2. 把oneflow_model.py中的layer尽可能的替换成LiBai中支持的layer, 只替换你想要的部分也可以(比如只替换Linear层), LiBai会自动把没有替换的layer 转换成分布式并行所需要的格式. 这一步是支持分布式推理的关键
3. 继承LiBai提供好的分布式推理的基类 basic.py ( https://github.com/Oneflow-Inc/libai/blob/main/libai/inference/basic.py ), 重载转换权重的函数, 按照pytorch那样写好预处理和后处理, 就可以进行分布式推理了.

下面链接里面有极其详细的步骤解答, 看在作者不仅授人以鱼, 还授人以渔, 还做了程序员最讨厌的文档活儿, 看官老爷们给LiBai个star吧.点了star就不会迷路了, 而且保不齐以后万一有个什么复现的任务, 这里面的知识点也用得上,至少可以用别人release出来的预训练权重来验证自己复现的model.py是否正确. 这叫什么, 这叫call back!

LiBai分布式推理介绍：https://github.com/Oneflow-Inc/libai/discussions/386

以MT5白嫖hugging face model为例子, 我们在2机4卡下面进行 模型并行2X流水并行2 的分布式推理, 跑起来的代码风格如下

    # test_inference.py
    from libai.inference.text_generation import TextGenerationPipeline
    from libai.utils import distributed as dist
    
    if __name__ == "__main__":
      pipeline = TextGenerationPipeline(
          "projects/MT5/configs/t5_inference.py",
          data_parallel=1,
          tensor_parallel=2,
          pipeline_parallel=2,
          pipeline_stage_id=[0] * 12 + [1] * 12,
          pipeline_num_layers=12 * 2,
          model_path="data_test/t5_inference_model",
          mode="huggingface",
      )
    
      text = ["summarize: She is a student, She is tall, She loves study"]
      dict1 = pipeline(text)
      if dist.is_main_process():
          print(dict1)

那么多机多卡的分布式推理脚本

在node0上输入指令:

NODE=2 NODE_RANK=0 ADDR=192.168.0.1 PORT=12345 bash tools/infer.sh test_inference.py 2

在node1上输入指令:

NODE=2 NODE_RANK=1 ADDR=192.168.0.1 PORT=12345 bash tools/infer.sh test_inference.py 2

细心的朋友已经发现了, LiBai下面可以通过设置pipeline_stage_id, 来让用户自己设置每个stage上group的层数是多少,方便在某些极端情况下(比如你的机器0很牛逼, 但是机器1很拉胯, 或者你的encoder计算量巨大,但是decoder计算量较小)手动实现负载均衡.

大模型训练

众所周知, 大家都喜欢做点"出格"的事情, 比如在上班的时候摸鱼, 在剧情片里面找爱情磕CP, 在爱情动作电影里面看剧情, 在vscode上面炒股, 在pxxxhub上面学微积分......

那么LiBai呢? 你甚至可以拿它来训练模型!!

介绍一下LiBai中目前支持的模型:

在 Projects：https://github.com/Oneflow-Inc/libai/tree/main/projects 下支持的模型:

下面来谈谈模型之外, LiBai有什么不一样的地方. 按照下面两段话术来说, 也就是核心竞争力在哪里

分布式配置和算法逻辑解耦

LiBai进行了模块化的设计, 使得分布式的配置和算法逻辑解耦, 这意味着什么? 这意味着用户只需要把大部分的注意力专注到算法逻辑上面, 而不用在苦恼怎么插入各种并行的代码了.

简单来说, 下面这些模块都可以在config.py中进行一键配置.

    # my_config.py
    from libai.config import get_config
    train = get_config("common/train.py").train
    optim = get_config("common/optim.py").optim
    graph = get_config("common/models/graph.py").graph
    
    # set dist
    train.dist.data_parallel_size = 2
    train.dist.tensor_parallel_size = 2
    train.dist.pipeline_parallel_size = 2
    # set model layers for pipeline
    train.dist.pipeline_num_layers = 24
    # set pipeline_stage_id according to your own needs.
    # if `None`, LiBai will use its own mode of distribution
    train.dist.custom_pipeline_stage_id = [0]*14 + [1]*10
    
    # set auto parallel in LiBai
    graph.auto_parallel.enabled = True
    
    # enable amp (fp16)
    train.amp.enabled = True 
    
    # enable gradient clipping
    optim.params.clip_grad_norm = 1.0
    optim.params.clip_grad_norm_type = 2.0
    
    # enable grad accumulation for 8 steps
    train.num_accumulation_steps = 8
    
    # enable activation checkpointing
    train.activation_checkpoint.enabled = True
    
    # enable zero for leval-2
    train.zero_optimization.enabled = True
    train.zero_optimization.stage = 2

单机和分布式代码几乎一致

下面给一个简单的2D并行(数据并行+模型并行)的MLP例子, 比如你的Linear层在16384这个维度上面比较大, 需要把它切分在不同的卡上才能装下, 那么在LiBai下面只需要如下所示就可以完成了, 几乎的单机的代码没有区别.

    from libai.layers.linear import Linear
    from oneflow import nn 
    
    # write a Simple 2D Parallel MLP
    class MLP_2D(nn.Module):
        def __init__(self,):
            super().__init__()
            self.linear1 = Linear(in_features=1024, out_features=16384, parallel="col")
            self.relu = nn.GELU()
            self.linear2 = Linear(in_features=16384, out_features=1024, parallel="row")
            self.dropout = nn.Dropout(p=0.5)
        
        def forward(self, x):
            x = self.linear1(x)
            x = self.relu(x)
            x = self.linear2(x)
            x = self.dropout(x)
            return x

支持一键转换onnx

和上次的年更大作一样, 本人对一键转onnx的执念可谓是相当之深了. 在LiBai下面我们同样以MT5为例子, 支持了一键转换onnx的功能, 点击以下的链接就可以体验到了.

MT5转onnx：https://github.com/Oneflow-Inc/libai/tree/main/libai/onnx_export

更详细的说明和教程会在LiBai中持续发布, 只要点一下收藏(也就是github上面的star按钮), 就可以持续跟进最新发展了.

结语

感谢各位看官老爷的阅读~, 如果这篇文章有让身为乐子人的读者快乐或者有启发的话, 请不要吝惜手中的收藏按钮, 欢迎去项目上面star, fork, watch三连.

GitHub -Oneflow-Inc/libai: LiBai(李白): A Toolbox for Large-Scale Distributed ParallelTraining：https://github.com/Oneflow-Inc/libai

也可以扫码加群~

引用

Detectron2：https://github.com/facebookresearch/detectron2

oneflow：https://github.com/Oneflow-Inc/oneflow

oneflow_convert：https://github.com/Oneflow-Inc/oneflow_convert

Megatron：https://github.com/NVIDIA/Megatron-LM

外记

鸣谢

之所以把这个放到最后, 是因为这个部分有点长, 可以选择性跳过, 但我会尽量把人员写全一些, 会有一些小伙伴支持多线开发, 但是为了避免名字重复,就只写一遍吧.

因为最近经济下行, 传递给每个人的不应该只有寒气, 还有来自于同伴的鼓励和认可.

LiBai框架开发: 廖星宇, 程鹏, 任天和, 谢子鹏, @梁德澎：https://www.zhihu.com/people/2a31b6e715c7ce64c917bc09f7c3b6f5 , 党凯, 黄威, 王义, 邵世通, 迟子秋, 刘驰, 张广军, 陈巧玲, 钟珊珊, 庄宇林

底层优化(oneflow): 成诚, 李新奇, 许啸宇, 张文骁, @BBuf：https://www.zhihu.com/people/f57573ed58d7168ced17f0ac662711b9 , 王迎港, @大缺弦：https://www.zhihu.com/people/33295f2791c588f9df071dddb701278a , 柳俊丞, @我不是zzk：https://www.zhihu.com/people/236dedf0a4bc04040b68aea7e2e6e06a , 郭冉,刘沛宏, 蔡晟航, 李一鹏, 谢暄, 陈后江, 李响, @Lyon：https://www.zhihu.com/people/38495a9b1f83e4a612543ce6c523df7f , 韩彬彬.

测试: 欧阳宇, 徐雍宁

宣传与翻译: @姚迟：https://www.zhihu.com/people/4ce7df44a2b6f1a9ca4233250eefd614 ,王金许, 上官士源, 董文文, 张雨珊