最全攻略：利用LightSeq加速你的深度学习模型

前言

LightSeq是字节跳动火山翻译团队开源的一款Transformer系列模型加速引擎，分为训练和推理两个部分。其中推理加速引擎早在2019年12月就已经开源，而训练加速引擎也在2021年6月开源。

项目地址：

https://github.com/bytedance/lightseq

LightSeq主要采用了CUDA算子融合、显存优化、参数连续化、层级式解码策略等技术，感兴趣的小伙伴可以阅读此前的文章：

训练引擎：

支持Transformer全流程训练加速，最高加速3倍！字节跳动LightSeq上新

推理引擎：

速度超快！字节跳动开源序列推理引擎LightSeq

本文详细讲解一下如何使用LightSeq来改造你的PyTorch模型，实现1.5-3倍的训练加速和5-10倍的推理加速。至于TensorFlow模型的加速，目前也已经支持，这里不会详细讲解，可以参考下面NeurST的代码：

https://github.com/bytedance/neurst/tree/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-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层。替换过程分为三个步骤：

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.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的命令不同的地方有这么几个：

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.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 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进行加速。