前言

LightSeq是字节跳动火山翻译团队开源的一款Transformer系列模型加速引擎,分为训练和推理两个部分。其中推理加速引擎早在2019年12月就已经开源,而训练加速引擎也在2021年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进行加速的整体流程依次为:

  1. 接入训练引擎进行模型训练,并保存模型参数。
  2. 加载模型参数,使用训练引擎的前向传播部分进行模型推理。
  3. 为了更快的推理速度,还可以将模型参数导出为protobuf或者hdf5格式。
  4. 使用推理引擎解析第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层为例,主要分为两个步骤:

  1. 使用LSTransformerEncoderLayer.get_config函数新建config。
  2. 新建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-basetransformer-bigbert-basebert-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层。替换过程分为三个步骤:

  1. 使用LSTransformerEncoderLayer.get_config函数新建config。
  2. 获取Hugging Face预训练好的BERT参数。
  3. 新建LightSeq的encoder层,即LSTransformerEncoderLayer类,使用config和预训练好的参数来初始化。

新建encoder层代码参见上一小节。注意在Hugging Face这个例子里,额外给LSTransformerEncoderLayer封装了一层LSHFTransformerEncoderLayer,主要是为了兼容原来的encoder输入形状。

示例代码在examples/training/huggingface中,运行sh run_glue.shsh 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目录存放着三个启动入口:trainvalidategeneratefs_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的命令不同的地方有这么几个:

  1. 启动入口从fairseq-train替换为了lightseq-train,这是因为在根目录setup.py里封装了--user-dir用户模块目录。如果还想继续用fairseq-train的话,就需要手动指定--user-dir fs_modules参数。
  2. 模型结构--arch需要在原来的基础上加上前缀ls_,用来指定使用LightSeq提供的Transformer模型。
  3. 优化器--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加速训练的预训练模型参数,所以导出更为复杂一些。

模型导出的核心思想就是:

  1. 首先创建一个protobuf对象Transformer或者hdf5的文件对象。
  2. 然后在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.pypython hf_bart_export.pypython 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 lsi
model = 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加速你的深度学习模型,最佳方式无外乎三步:

  1. 接入LightSeq训练引擎的模型组件,构建模型,进行训练,保存checkpoint。
  2. 将checkpoint转换为protobuf或者hdf5格式,LightSeq的组件可以调用现成的转换接口,其它的需要自己手写转换规则。
  3. 调用LightSeq推理引擎,加载上一步中导出的模型,进行快速推理。

目前LightSeq已经被广泛应用在字节跳动公司内外各项业务和学术研究上,支持了标准的Transformer、BERT、BART、GPT2、ViT等多种Transformer系列模型。只要你的模型中包含有Transformer的部分组件,例如encoder层,就可以直接调用LightSeq进行加速。

联系方式

如果在使用中遇到任何问题,或者有任何需求和建议,都可以在github issue中提出,或者加入LightSeq的飞书用户群:

AppLink

Logo

更多推荐