机器之心专栏

机器之心编辑部

Transformer已经成为众多NLP任务以及部分CV任务的主流模型，但由于硬件资源匮乏，很多高校实验室或者公司都无法训练很大的模型，而降低批处理大小等措施又会导致训练时间成倍增加。针对这一痛点，字节跳动推出了LightSeq训练加速引擎，对Transformer训练的整个计算过程进行了优化，最多可以实现3倍以上的加速。

如今，NLP和CV领域的大部分任务都将Transformer作为基础模型。而早在年12月，字节跳动就开源过一款Transformer类模型推理加速引擎——LightSeq[1]。作为业界第一款支持多种模型和解码方法的推理加速引擎，LightSeq的推理速度快于其它同类软件，更是远远超过了TensorFlow和PyTorch。

今天，LightSeq的最新版本发布了！

这次发布引入了引擎方面的重大更新——支持了Transformer全流程训练加速，在不同的批处理大小下相比主流训练库最高可加速3倍多！至此从训练到推理部署的整个过程都已被LightSeq打通。

那么它到底采用了哪些技术呢？笔者将根据LightSeq公布的资料为你逐个揭晓。

LightSeq能做什么？

图1：Transformer模型结构图（以机器翻译为例）

自年被谷歌提出之后，Transformer模型[2]成为了众多NLP任务以及部分CV任务的主流模型，尤其是机器翻译、文本生成、文本摘要、时间序列预测等任务。图1是机器翻译任务使用Transformer进行训练的一个例子。

但由于硬件资源匮乏，很多高校实验室或者公司都无法训练很大的模型，而降低批处理大小等措施又会导致训练时间成倍增加。因此，如何利用有限的资源尽快训练出模型成为了亟待解决的问题。

针对这一痛点，字节跳动推出了LightSeq训练加速引擎，对Transformer训练的整个计算过程进行了优化。从词嵌入层、编码层、解码层，到最后的损失函数层，从前向传播、反向传播、梯度同步，到最后的参数更新，LightSeq都进行了细致的性能分析和优化。

以机器翻译任务为例，在该任务中，我们仅需要修改几行代码开启LightSeq训练加速，就可以实现最多3倍以上的加速。

总的来说，LightSeq具有如下几个优点：

1.支持Transformer的完整加速。

LightSeq是业界第一款完整支持整个Transformer模型加速的训练引擎，包括了词嵌入层、编码层、解码层、损失函数层等高效自定义层。相比之下，另一款知名的深度学习优化引擎DeepSpeed[3]仅支持编码层的加速，因此只能用在BERT等模型上，局限性较大。

2.训练速度快。

LightSeq训练速度非常快。例如在WMT14[4]英德机器翻译任务上，利用英伟达最新的训练显卡A，相比于主流序列生成库，LightSeq最快仅需要三分之一的训练时间。

3.功能全面，简单易用。

LightSeq提供了高效的TensorFlow和PyTorch自定义层供用户灵活使用，可以自由插入到HuggingFace等主流训练库中。除此之外，它还和当前流行的训练库如Fairseq[5]、NeurST[6]等做了深度集成。用户在安装LightSeq后，只需要修改几个命令行参数，就能在这些训练库上使用LightSeq。

4.提供丰富的二次开发工具。

LightSeq提供了完整的CUDAkernel和Transformer自定义层的单元测试功能，可以测试自定义算子的正确性，同时分析出运行时间和加速比，帮助开发者更快地验证功能的正确性和有效性。

下表1列举了不同训练加速引擎之间支持的功能对比：

表1：不同训练加速引擎支持功能对比

如何快速上手LightSeq？

LightSeq提供了多种便捷的运行方式，点击文末链接可以快速体验。这里简单讲解一下快速接入LightSeq进行训练的两种方式。

使用自定义层

LightSeq提供了许多自定义层供用户灵活使用。

例如，如果你想将Transformer模型中的编码层替换为LightSeq的编码层，只需要提供一个编码层参数，用来初始化LightSeq编码层，然后就可以用它来替换原始的编码层，加速模型训练。详细代码如下（这里省略了部分配置参数）：

fromlightseq.training.ops.pytorch.transformer_encoder_layerimportLSTransformerEncoderLayer

config=LSTransformerEncoderLayer.get_config(

max_batch_tokens=,

max_seq_len=,

...

)

enc_layer=LSTransformerEncoderLayer(config)

一键启动

LightSeq还和几个当前流行的训练库（例如Fairseq和NeurST）做了深度集成，仅需修改很少的代码就能开启LightSeq加速。

Fairseq

LightSeq为Fairseq提供了一套完整便捷的Transformer训练样例。

首先，你需要安装Fairseq以及必要的第三方库，然后用如下命令安装LightSeq加速库：

pipinstalllightseq

接着就可以通过LightSeq提供的启动器，灵活地指定使用LightSeq优化版本的Transformer模型、参数优化器和损失函数。下面是启动参数样例（这里省略了和LightSeq无关的参数）：

lightseq-trainDATA_PATH\

--archls_transformer_wmt_en_de_big_t2t\

--optimizerls_adam\

--criterionls_label_smoothed_cross_entropy

NeurST

NeurST是一款同时支持TensorFlow和PyTorch的开源序列生成库，可以用来做文本生成、机器翻译和语音翻译。LightSeq已经与其进行了深度融合，无需用户修改代码和启动参数即可直接进行加速训练。

首先需要安装NeurST，官方提供了详细的安装教程[6]。

然后安装TensorFlow版本的LightSeq，命令如下：

pipinstall

转载请注明:http://www.aideyishus.com/lkjg/3367.html