只用几行代码，我让模型『训练』加速了3倍以上！

编程算法

发布日期: 2021-06-21

作者: 韦阳

文章字数: 1.2k

阅读时长: 4 分

阅读次数:

最近总有学妹问我，论文要截稿了，模型来不及跑怎么办？

还有学妹问我，有好多idea，验证一个就要跑一周怎么办？

这时候我想起了下面这张图，我想这句话反映了大多数从事人工智能的科研工作者们目前的状态。

于是我告诉学妹，想要快，找我就对了，我教你怎样让你的模型训练加速3倍以上！

这里我们需要用到的就是字节跳动AI Lab最近开源的训练加速引擎——LightSeq。

具体的原理这里就不做过多介绍了，另一篇文章里写的很详细：训练加速3倍！字节跳动推出业界首个NLP模型全流程加速引擎。今天我来教大家如何使用LightSeq。

安装步骤

LightSeq依赖第三方库Ninja对CUDA进行编译，所以先装它：

pip install ninja

源码安装

你可以从源码进行安装，使用如下命令：

git clone https://github.com/bytedance/lightseq.git
cd lightseq
pip install -e .

如果你想执行LightSeq提供的现成样例，或者使用它的单元测试工具，那最好从源码安装。

pip安装

当然如果你想直接调用LightSeq的接口，不需要它的样例或者单元测试工具，我更推荐你用下面pip的方式安装，更加方便：

pip install lightseq

使用教程

这次讲解的所有样例都放在了lightseq/training/examples目录下，推荐大家采用上面源码安装的方式安装LightSeq，这样可以直接运行样例。

Hugging Face

Hugging Face是目前用的最多的预训练模型库了吧，主要是用起来太方便了，模型也很全。直接pip install transformers安装即可。

以BERT在NER任务上微调为例，直接运行LightSeq提供的脚本就行：

sh lightseq/training/examples/huggingface/run_ner.sh

Fairseq

Fairseq是目前最主流的序列生成库之一，用来做机器翻译、文本生成等任务都是非常方便的。安装的话也很简单，直接pip install fairseq即可。

LightSeq同样提供了现成的运行脚本，如果想运行LightSeq加速后的模型，执行如下命令：

sh lightseq/training/examples/fairseq/ls_fairseq_wmt14en2de.sh

再来看看细节，一般如果我们想用Fairseq来训练一个机器翻译模型，通常首先会准备好数据集，然后执行如下命令：

fairseq-train DATA_DIR \
    --arch transformer_wmt_en_de_big_t2t \
    --optimizer adam \
    --criterion label_smoothed_cross_entropy \
    ...

注意这里我们只列出了同LightSeq有关的三个参数：--arch、--optimizer和--criterion，分别指定了模型结构、参数优化器和损失函数。

如果想用LightSeq进行加速，直接将上面的运行命令改为下面这样：

lightseq-train DATA_DIR \
    --arch ls_transformer_wmt_en_de_big_t2t \
    --optimizer ls_adam \
    --criterion ls_label_smoothed_cross_entropy \
    ...

注意改动的地方有4个。fairseq-train改成lightseq-train，这是为了导入LightSeq的目录。--arch、--optimizer和--criterion都加上了ls_前缀，这样就快速替换为了LightSeq的组件。

自定义模型

当然绝大多数时候你可能只是想自己搭建一个Transformer模型，然后用来做自己的事情。

以编码层为例，你只需要为每一层创建一个配置对象，然后用它创建LightSeq的编码层即可。

我写了一份完整的训练代码，非常浅显易懂，看注释就行了：

import torch
from lightseq.training.ops.pytorch.transformer_encoder_layer import LSTransformerEncoderLayer

def train(model, inputs, masks):
    inputs = inputs.to(device="cuda:0")
    masks = masks.to(device="cuda:0")
    model.to(device="cuda:0")
    model.train()
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)

    for epoch in range(1000):
        opt.zero_grad()
        outputs = model(inputs, masks)
        loss = torch.square(outputs).mean()
        loss.backward()
        opt.step()
        if epoch % 200 == 0:
            print("epoch {:>3d}: loss = {:>5.3f}".format(epoch, loss))

if __name__ == "__main__":
    # 定义LightSeq配置
    config = LSTransformerEncoderLayer.get_config(
        max_batch_tokens=4096,
        max_seq_len=256,
        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=True,
        fp16=False,
        local_rank=0
    )

    # 随机生成输入
    bsz, sl = 10, 80
    inputs = torch.randn(bsz, sl, config.hidden_size)
    masks = torch.zeros(bsz, sl)

    # 定义LightSeq编码层并进行训练
    model = LSTransformerEncoderLayer(config)
    train(model, inputs, masks)