前言
LightSeq是字节跳动火山翻译团队开源的一款Transformer系列模型加速引擎,分为训练和推理两个部分。其中推理加速引擎早在12月就已经开源,而训练加速引擎也在6月开源。
项目地址:
GitHub - bytedance/lightseq: LightSeq: A High Performance Library for Sequence Processing and Generation
LightSeq主要采用了CUDA算子融合、显存优化、参数连续化、层级式解码策略等技术,感兴趣的小伙伴可以阅读此前的文章:
训练引擎:
godweiyang:训练加速3倍!字节跳动推出业界首个NLP模型全流程加速引擎
推理引擎:
机器之心:速度超快!字节跳动开源序列推理引擎LightSeq
本文详细讲解一下如何使用LightSeq来改造你的PyTorch模型,实现1.5-3倍的训练加速和5-10倍的推理加速。至于TensorFlow模型的加速,目前也已经支持,这里不会详细讲解,可以参考下面NeurST的代码:
GitHub - bytedance/neurst at lightseq
整体流程
使用LightSeq进行加速的整体流程依次为:
接入训练引擎进行模型训练,并保存模型参数。加载模型参数,使用训练引擎的前向传播部分进行模型推理。为了更快的推理速度,还可以将模型参数导出为protobuf或者hdf5格式。使用推理引擎解析第3步中导出的模型,并进行模型推理。
模型训练
LightSeq提供了封装好的embedding、encoder、decoder、cross entropy和adam类,可以接入到你自己的模型中替换原有的模型。
LightSeq还提供了现成的Fairseq、Hugging Face、DeepSpeed DeepSpeed可以用于大规模训练Speed、NeurST等样例。如果你用这几个训练库的话,就可以直接使用。如果你是自己的模型,那也可以手动接入LightSeq。这几个样例代码都在examples/training
目录下。
自定义模型
首先引入所有可能用到的头文件:
from lightseq.training import (LSTransformer,LSTransformerEmbeddingLayer,LSTransformerEncoderLayer,LSTransformerDecoderLayer,LSCrossEntropyLayer,LSAdam,)
以新建encoder层为例,主要分为两个步骤:
使用LSTransformerEncoderLayer.get_config
函数新建config。新建LightSeq的encoder层,即LSTransformerEncoderLayer
类,使用config来初始化。
一个典型的例子如下:
config = LSTransformerEncoderLayer.get_config(model="bert-base",max_batch_tokens=4096,max_seq_len=512,fp16=True,local_rank=0,)layer = LSTransformerEncoderLayer(config)
其中max_batch_tokens
指定了训练过程中一个batch最大可能的单词数,max_seq_len
指定了句子的最长长度。model
提供了四种现成的模型配置:transformer-base
、transformer-big
、bert-base
和bert-big
。
当然如果你想用自己的模型配置,也可以手动补全所有的参数:
config = LSTransformerEncoderLayer.get_config(max_batch_tokens=4096,max_seq_len=512,hidden_size=1024,intermediate_size=4096,nhead=16,attn_prob_dropout_ratio=0.1,activation_dropout_ratio=0.1,hidden_dropout_ratio=0.1,pre_layer_norm=False,activation_fn="gelu",fp16=True,local_rank=0,)layer = LSTransformerEncoderLayer(config)
除了encoder以外,embedding、decoder、cross entropy和adam也可以用同样的方法新建,最后和你自己写的模型一样进行训练即可。
此外LightSeq还提供了完整的Transformer类LSTransformer
,可以直接新建一整个Transformer:
config = LSTransformer.get_config(model="transformer-base",max_batch_tokens=4096,max_seq_len=512,vocab_size=32000,padding_idx=0,num_encoder_layer=6,num_decoder_layer=6,fp16=True,local_rank=0,)model = LSTransformer(config)
示例代码在examples/training/custom
中,可以直接运行python run.py
查看效果。
Hugging Face
以Hugging Face官方提供的run_glue.py
为例,一般首先都是用AutoModel.from_pretrained
函数新建模型model,然后进行训练。
为了接入LightSeq,需要将model中的所有encoder层替换为LightSeq版本的encoder层。替换过程分为三个步骤:
使用LSTransformerEncoderLayer.get_config
函数新建config。获取Hugging Face预训练好的BERT参数。新建LightSeq的encoder层,即LSTransformerEncoderLayer
类,使用config和预训练好的参数来初始化。
新建encoder层代码参见上一小节。注意在Hugging Face这个例子里,额外给LSTransformerEncoderLayer
封装了一层LSHFTransformerEncoderLayer
,主要是为了兼容原来的encoder输入形状。
示例代码在examples/training/huggingface
中,运行sh run_glue.sh
和sh run_ner.sh
分别可以查看LightSeq在GLUE和NER任务上的加速效果。
注意Hugging Face BERT的fine-tune任务很不稳定,经常会不收敛,这时候可以尝试修改运行脚本中的--seed
参数。
Fairseq
Fairseq主要用于一些生成任务,使用LightSeq加速的原理是一样的,都是需要将各自组件替换为LightSeq对应的组件。
LightSeq对Fairseq做了非常完整的替换,将embedding、encoder、decoder、cross entropy和adam全部替换为了LightSeq对应的部分,来达到极致的加速效果。
示例代码在examples/training/fairseq
目录下,其中fs_cli
目录存放着三个启动入口:train
、validate
和generate
,fs_modules
目录存放着用LightSeq封装好的几个Transformer组件。
直接运行sh ls_fairseq_wmt14en2de.sh
即可自动下载数据并运行WMT14英德机器翻译任务。脚本中主要的运行命令如下:
lightseq-train /tmp/wmt14_en_de/ \--task translation \--arch ls_transformer_wmt_en_de_big_t2t --share-decoder-input-output-embed \--optimizer ls_adam --adam-betas '(0.9, 0.98)' --clip-norm 0.0 \--lr 5e-4 --lr-scheduler inverse_sqrt --warmup-updates 4000 --weight-decay 0.0001 \--criterion ls_label_smoothed_cross_entropy --label-smoothing 0.1 \--max-tokens 8192 \--eval-bleu --eval-bleu-args '{"beam": 5, "max_len_a": 1.2, "max_len_b": 10}' \--eval-bleu-detok moses --eval-bleu-remove-bpe --eval-bleu-print-samples \--best-checkpoint-metric bleu \--maximize-best-checkpoint-metric --fp16
注意到和一般运行Fairseq的命令不同的地方有这么几个:
启动入口从fairseq-train
替换为了lightseq-train
,这是因为在根目录setup.py
里封装了--user-dir
用户模块目录。如果还想继续用fairseq-train
的话,就需要手动指定--user-dir fs_modules
参数。模型结构--arch
需要在原来的基础上加上前缀ls_
,用来指定使用LightSeq提供的Transformer模型。优化器--optimizer
和损失函数--criterion
都需要在原来的基础上加上前缀ls_
,指定使用LightSeq对应的组件。
DeepSpeed
DeepSpeed主要用于大规模训练,也提供了Transformer的encoder层CUDA实现,不过效率没有LightSeq高。
LightSeq提供了Fairseq+DeepSpeed分布式训练的使用样例,将启动器替换成了deepspeed
,手动指定--user-dir
目录,还需要指定DeepSpeed的配置文件deepspeed_config
,其它参数和上一节Fairseq样例一模一样。
使用时运行sh ds_fairseq_wmt14en2de.sh
即可,和上一小节一样都是用Fairseq运行WMT14英德机器翻译任务。
模型导出
在模型训练完之后,直接load保存的checkpoint就可以继续fine-tune或者推理。但是这样调用的是训练引擎的推理部分,也就是模型的前向传播。这部分代码需要频繁在python和c++之间切换,并且前向过程中计算了很多反向传播才需要用到的变量。因此速度不如纯粹的推理引擎快。
而要想使用LightSeq的推理引擎,就必须先将checkpoint转变为protobuf或者hdf5的格式。
LightSeq提供了每个组件的导出接口,如果你使用了LightSeq的模型组件,那么导出将变得非常容易。只需要引入下面的头文件即可:
from lightseq.training import (export_ls_config,export_ls_embedding,export_ls_encoder,export_ls_decoder,)
这四个函数分别可以导出推理引擎所需要的配置信息、embedding参数、encoder参数和decoder参数。而如果有其他部分的参数没包括在这里面(例如输出到词表的映射矩阵),则需要自己进行导出,详见下面的教程。
LightSeq对Hugging Face的BERT、BART、GPT2三种模型,以及Fairseq+LightSeq、LightSeq的Transformer模型都提供了模型导出的样例,代码在examples/inference/python/export
目录下。其中Hugging Face的模型都是没有采用LightSeq加速训练的预训练模型参数,所以导出更为复杂一些。
模型导出的核心思想就是:
首先创建一个protobuf对象Transformer
或者hdf5的文件对象。然后在checkpoint中提取出参数值,将其赋值给Transformer
或者hdf5文件对象中对应的参数。
这个过程麻烦的就是提取并且对应赋值的过程,LightSeq提供了一系列方便的操作函数。
Fairseq
执行python ls_fs_transformer.py
可以导出上一章节中Fairseq+LightSeq训练样例得到的模型。
以protobuf导出为例,观察代码可以看到主体部分如下(省略了部分参数):
file = Transformer()encoder_state_dict, decoder_state_dict = _extract_weight(state_dict)export_ls_embedding(file, encoder_state_dict, is_encoder=True)export_ls_embedding(file, encoder_state_dict, is_encoder=False)export_ls_encoder(file, encoder_state_dict)export_ls_decoder(file, decoder_state_dict)export_fs_weights(file, state_dict)export_ls_config(file)
首先需要用户自己将state_dict拆分成encoder和decoder两部分,这主要是因为设计时考虑到有些用户只会用到encoder的导出(例如BERT)。并且LightSeq无法知道用户模型的最外层参数名叫啥,万一不叫encoder,而叫enc之类的呢?所以交给用户自己拆分更加合理。
然后分别导出encoder的embedding、decoder的embedding、encoder和decoder参数,这几部分都直接调用LightSeq提供的接口就行了。LightSeq会自动帮你把解析出来的参数导出到定义的Transformer
类里。
接着需要处理一下Fairseq中与LightSeq无关的一些参数,例如encoder和decoder的layer norm参数等等。export_fs_weights
函数需要用户自己实现,核心思想就是找到state_dict中的参数名,将其赋值给Transformer
类里对应的变量就行了。
最后设置一下Transformer
类里所有的配置参数就行了。
hdf5的用法类似,LightSeq都将其封装在同样的函数里了,只需要指定save_pb=False
即可。
Hugging Face
执行python hf_bert_export.py
、python hf_bart_export.py
和python hf_gpt2_export.py
三个文件分别可以导出BERT、BART和GPT2的预训练模型。
因为Hugging Face的模型参数都是预训练得到的,所以LightSeq无法识别参数名是什么样的,只能用户自己编写导出规则,具体参考上面三个导出样例即可。
LightSeq Transformer
使用LightSeq提供的Transformer进行训练的话,参数名LightSeq都知道的一清二楚,因此可以直接使用LightSeq提供的导出接口进行转换。过程和上面的Fairseq+LightSeq类似。
具体样例可以执行python ls_transformer_export.py
,同时得到protobuf和hdf5格式的模型导出文件,并且对比两者生成的结果。这里的checkpoint可以使用上一章节中自定义模型
小节中训练得到的模型。
自定义模型
因为自定义的模型参数LightSeq无法识别参数名,所以需要用户自己编写转换规则。
举一个简单的例子,假设用户模型中有个encoder的输出部分的layer norm参数,state_dict中的参数名叫做encoder.layer_norm.weight
。那么可以按如下方式进行转换:
transformer = Transformer()enc_norm_w = state_dict["encoder.layer_norm.weight"].flatten().tolist()transformer.src_embedding.norm_scale[:] = enc_norm_w
模型推理
得到导出的protobuf或者hdf5模型后,推理就变得十分简单,核心代码就三行:
import lightseq.inference as lsimodel = lsi.Transformer("transformer.pb", 8)output = model.infer([[1, 2, 3], [4, 5, 6]])
首先定义一个Transformer
类用来加载模型参数,指定load的protobuf文路径和batch_size大小。
然后调用infer
函数进行推理,传入的输入参数必须是list或者numpy类型,且必须是二维。
LightSeq在examples/inference/python/test
目录下提供了三个Hugging Face模型推理的样例,此外上一小节中examples/inference/python/export
中的ls_transformer_export.py
代码也包含了导出后推理的过程。
最佳实践
总结一下,使用LightSeq加速你的深度学习模型,最佳方式无外乎三步:
接入LightSeq训练引擎的模型组件,构建模型,进行训练,保存checkpoint。将checkpoint转换为protobuf或者hdf5格式,LightSeq的组件可以调用现成的转换接口,其它的需要自己手写转换规则。调用LightSeq推理引擎,加载上一步中导出的模型,进行快速推理。
目前LightSeq已经被广泛应用在字节跳动公司内外各项业务和学术研究上,支持了标准的Transformer、BERT、BART、GPT2、ViT等多种Transformer系列模型。只要你的模型中包含有Transformer的部分组件,例如encoder层,就可以直接调用LightSeq进行加速。
联系方式
如果在使用中遇到任何问题,或者有任何需求和建议,都可以在github issue中提出,或者加入LightSeq的飞书用户群:
AppLink