李白：你的模型权重很不错，可惜被我没收了

撰文｜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的分布式推理，跑起来的代码风格如下：

# test_inference.py
fromlibai.inference.text_generationimport TextGenerationPipeline

fromlibai.utilsimport 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计算量较小）手动实现负载均衡。

大模型训练

众所周知，大家都喜欢做点"出格"的事情，比如在上班的时候摸鱼，在VScode上面炒股......

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

‍

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

下面来谈谈模型之外，LiBai有什么不一样的地方，换句话说，也就是核心竞争力在哪里？

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

LiBai进行了模块化的设计，使得分布式的配置和算法逻辑解耦，这意味着什么?

这意味着用户只需要把大部分的注意力专注到算法逻辑上面，而不用再苦恼怎么插入各种并行的代码了。

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

# my_config.py
fromlibai.configimport 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下面只需要如下所示就可以完成了，几乎跟单机代码没有区别。

fromlibai.layers.linearimport Linear

fromoneflowimport nn 


# write a Simple 2D Parallel MLP
classMLP_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)


defforward(
self, x):

 x 
=self.linear1(x)

 x 
=self.relu(x)

 x 
=self.linear2(x)

 x 
=self.dropout(x)

return x

支持一键转换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

继续阅读

阅读原文