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

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

还有学妹问我,有好多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)

训练速度

说了这么多,实际速度到底怎么样?我用Fairseq测了一下训练的总耗时:

不同模型大小、不同批处理大小、不同显卡上加速效果都是有区别的,但总体上都能缩短一半左右的训练时间。

如果你的显卡比较老旧(我相信大多数学校实验室都是这样的),显存又比较小,那么批处理大小只能设置的很小,那加速比甚至能达到3倍以上。

项目地址

学妹试了直叫好,让我下次别再这么快了。

你也别忘了点个star,让更多的人享受到极速的快乐。

LightSeq地址:
https://github.com/bytedance/lightseq


   转载规则


《只用几行代码,我让模型『训练』加速了3倍以上!》 韦阳 采用 知识共享署名 4.0 国际许可协议 进行许可。
 上一篇
网传字节跳动实习生删除所有轻量级模型,还有救吗? 网传字节跳动实习生删除所有轻量级模型,还有救吗?
这两天字节跳动的一个新闻上了热搜,相信大多数人应该都已经听说了: 事情大致经过就是字节跳动一个新实习生想删除HDFS下一个目录,可能是存储空间不够了吧。然后经过多方求证,确定了这个目录已经没用了,接着删除了,还使用了无法恢复的命令删除的。
2021-06-25
下一篇 
训练加速3倍!字节跳动推出业界首个NLP模型全流程加速引擎 训练加速3倍!字节跳动推出业界首个NLP模型全流程加速引擎
如今NLP和CV领域的大部分任务都将Transformer作为基础模型。而早在2019年12月,字节跳动曾经开源过一款Transformer类模型推理加速引擎——LightSeq。作为业界第一款支持多种模型和解码方法的推理加速引擎,Lig
2021-06-21
  目录