You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
# test_inference.pyfromlibai.inference.text_generationimportTextGenerationPipelinefromlibai.utilsimportdistributedasdistif__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)
ifdist.is_main_process():
print(dict1)
reacted with thumbs up emoji reacted with thumbs down emoji reacted with laugh emoji reacted with hooray emoji reacted with confused emoji reacted with heart emoji reacted with rocket emoji reacted with eyes emoji
-
前言
LiBai
发布新模块, 分布式推理适用场景:
GPU
卡无法容纳一个大模型的权重时, 需要把权重切分到多个卡上, 来完成模型的推理支持的功能:
layer
层数, 方便进行负载均衡libai
的model
, 和torch
的model
进行分布式推理(需要重载加载模型的模块)demo
的演示一份简单的多机多卡分布式(模型并行2X流水并行2)运行代码
多机多卡的分布式推理脚本
在
node0
上输入指令:在node1上输入指令:
如何使用
pytorch
的模型进行分布式推理?在
LiBai
下面加载torch
的模型进行分布式推理之前, 我们需要有一些预备知识预备知识
一个完整的模型由两个部分构成:
model.py
model_best.pth
那么假设我们在框架A下面, 有
modelA.py
和model_best_A.pth
, 我们想在框架B上面跑起来这个框架A下面的模型, 应该怎么做呢?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
模型的分布式推理在有了预备知识了以后, 再来看看怎么在
LiBai
下面跑pytorch
模型的分布式推理.主要分为以下的几步:
torch
的算子替换为oneflow
: 把torch_model.py
下面的torch
全部替换为oneflow
, 得到oneflow_model.py
.oneflow_model.py
中的layer
尽可能的替换成LiBai
中支持的layer
.model_config.py
, 这个部分比较简单, 如果没有完整训练的train_config.py
也没有关系, 只用写一份调用model的config就可以了, 可以参考dalle2_config.pylibai/inference/basic.py
, 写好预处理和后处理, 把第三步写好的转换权重的脚本重载到方法load_pretrain_weight
中.步骤1:把
torch
的算子替换为oneflow
得益于
oneflow
和pytorch
的api
高度一致, 通常情况下, 可以直接把torch
代码中的import torch
, 直接替换成import oneflow
, 正常情况下面, 可以把torch_model.py
下面所有的torch
关键字全部替换为oneflow
.在成功跑通了以后, 我们就有了一份以
oneflow
为底层算子的单机代码oneflow_model.py
了.如果没有办法跑通, 可以在报错的地方查看一下, 可能是算子没有对齐, 这个时候可以去
oneflow
下提issue
, 或者直接换另外一种不影响运行结果的写法直接绕过去.步骤2: 替换为LiBai中的layer
在
LiBai
下面提供了封装好的layers
可以进行调用, 这些layers
是基于oneflow
的再次封装, 目的是在于用户可以直接调用这些封装好的模块, 而不用自己去定义sbp
和placement
, 这些layer
是提供分布式推理的关键.拿一个最常用也最实用的例子: 在
transformer
的模型中, 一般Linear
层是运用最多的层, 也是最适合切分其权重到不同的gpu
上面, 以完成分布式训练或者推理的层.得益于
LiBai
的设计, 用户可以把模型中的layer
部分替换为LiBai
的layer
, 其他的代码可以不动. 所以对于用户来说, 如果Linear
层是单卡GPU
放不下的瓶颈, 那么用户可以只把代码中的Linear
层用LiBai
的Linear
层进行替换.其中
LiBai
的Linear
层提供了两种并行方式parallel="row"
和parallel="col"
, 分别代表着模型并行下, 参数是按照按行切分以及以及按列切分. 通常来说, 多个连续的Linear层, 按照col->row->col->row->col->....
切分方式的运算效率是比较高的.步骤3: 实现自己的
ModelLoader
用于加载模型权重利用训练好的模型权重是推理的关键之一,一般来说模型能够正确加载权重需要满足以下几点:
huggingface
中有丰富的模型权重资源,可以获取到本地。state_dict.keys()
与model.state_dict().keys()
一致。state_dict.values()
与model.state_dict().values()
一致。一般满足上述3点的
pytorch
模型即可直接加载模型权重,但是LiBai
中支持分布式推理功能,可以在多机多卡的形式下正确加载模型权重并且其中的tensor
都以分布式的形式存在,但这些处理细节不需要用户考虑,只需确保上述3点正确。确保权重文件中权重的
state_dict.keys()
与model.state_dict().keys()
一致:由于模型的定义风格不同,
LiBai
中的模型参数命名与Transformers
仓库有所不同,例如以下是LiBai
中和Transfomers
中Bert
的Embedding
层的命名差异,其中的vocab_embeddings
与word_embeddings
是等价的:所以在
ModelLoader
中需要实现_convert_state_dict
,用于将huggingface
获取的权重文件中的key
转为LiBai
的形式,可以参考BertLoader._convert_state_dict.另外在
ModelLoader
中还需要实现_load_config_from_json
,用于将huggingface
获取的config.py
中的模型配置加载到LiBai
中,可以参考BertLoader._load_config_from_json.确保权重文件中权重的
state_dict.values()
与model.state_dict().values()
一致经过上述的步骤后已经实现了key值正确匹配,最后还需确认values值是一致的,在LiBai的模型实现中有两点需要注意:
query,key,value
的计算方式:LiBai中的query,key,value
的计算方式是固定的,Transformers中query,key,value
的计算方式在不同的模型有所区别,所以当qkv的计算方式不同时需要使用_fix_qkv_ordering
方法改变LiBai中query,key,value
tensor的排列顺序来与huggingface保持一致,可以参考BertLoader
中的做法:here。layernorm层的放置位置可以通过
apply_residual_post_layernorm
参数来改变,其原理可以参考LiBai
的文档How to use Huggingface’s pretrained weights in LiBai.目前
LiBai
当中支持了Bert, Roberta, GPT2, T5, MT5, Swin, Swin2, Vit
模型加载huggingface
权重的方法,使用方式可以参考LiBai
文档,如果需要实现暂未支持的模型的权重加载也可以参考该文档实现:ModelLoaderHuggerFace.步骤4: 编写
config.py
如果使用的是
LiBai
训练好的model
进行分布式推理, 那么直接采用训练时的train_config.py
即可, 参考Couplets下的distribute_infer.py如果是迁移的
pytorch
的model
, 在这种情况下面, 没有完整训练的train_config.py
也没有关系, 只用写一份调用model的model_config.py
就可以了(记得要包含语句model=LazyCall(...)(...)
), 可以参考dalle2_config.py步骤5: 编写自己的
PipelineInference.py
在
LiBai
中已经提供好了基类libai/inference/basic.py, 用户只需要重载自己需要的函数就可以了. 可以参考text_generation.py如果是采用
LiBai
训练出来的model
, 那么形式就比较简单了, 可以直接参考Couplets下的distribute_infer.py简单说明一下在
LiBai
中需要重载的函数:步骤5: 进行分布式推理
在步骤4中编写好了自己的
PipelineInference.py
以后, 就可以进行分布式的推理了. 如前言的demo
演示中提到的但是值得注意的是, 如果是
load_pytorch
的model
, 由于我们只在模型初始化的时候替换了LiBai
的layers, 所以在这种情况下面, 只支持tensor_parallel
的并行方式.这是因为如果想要支持
pipeline_parallel
的话, 还需要在model.forward()
中加入语句x.to_global()
(参考LiBai_code.py), 把上一个pipeline_stage
的中间结果, 同步到本pipeline_stage
中来, 这个和hugging_face中的x.to(cuda_divce)
(参考hugging_face_code.py)原理是一致的.至于具体在哪个语句上面加上
to_global()
语句, 可以打开pipeline_parallel
运行一遍, 定位到报错的语句添加即可.Beta Was this translation helpful? Give feedback.
All reactions