LightSeq：面向序列处理与生成的高性能库

目录

• 版本说明
• 简介
  • 支持矩阵
• 性能表现
  • Transformer 训练加速
  • BERT 训练加速
  • Transformer 推理加速
  • BERT 推理加速
• 安装
  • 通过 PyPI 安装
  • 从源码构建
• 快速开始
  • 从零开始使用 LightSeq 训练
  • 基于 Fairseq 使用 LightSeq 训练
  • 基于 Hugging Face BERT 使用 LightSeq 训练
  • 基于 Fairseq 使用 LightSeq 推理
  • 基于 Hugging Face BERT 使用 LightSeq 推理
  • 使用推理服务器部署 LightSeq
• 引用我们
• 我们正在招聘！

版本说明

[2022.10.25] 发布 v3.0.0 版本，支持 int8 混合精度训练与推理。[中文介绍]

[2021.06.18] 发布 v2.0.0 版本，支持 fp16 混合精度训练。[中文介绍]

[2019.12.06] 发布 v1.0.0 版本，支持 fp16 混合精度推理。[中文介绍]

简介

LightSeq 是一个用 CUDA 实现的高性能序列处理与生成训练及推理库。它能够高效计算 BERT、GPT、Transformer 等现代 NLP 和 CV 模型，因此非常适用于机器翻译、文本生成、图像分类及其他序列相关任务。

该库基于 CUDA 官方库（cuBLAS、Thrust、CUB）以及为 Transformer 模型家族特别融合和优化的自定义核函数构建。除了模型组件外，推理库还提供了基于 TensorRT 推理服务器 的易于部署的模型管理和服务后端。使用 LightSeq，您可以轻松开发修改后的 Transformer 架构，只需少量额外代码。

LightSeq 的训练和推理速度极快。以下是整体性能表现：
* 与 PyTorch fp16 训练相比，LightSeq fp16 训练最高可加速 3 倍。
* 与 PyTorch QAT（即量化感知训练）相比，LightSeq int8 训练最高可加速 5 倍。
* 与 PyTorch fp16 推理相比，LightSeq fp16 和 int8 推理分别最高可加速 12 倍 和 15 倍。

支持矩阵

LightSeq 支持多种特性，如下表所示。
| 特性 | 支持列表 |
| -------------- | ---------------------------------------------------------------- |
| 模型 | Transformer, BERT, BART, GPT2, ViT, T5, MT5, XGLM, VAE, 多语言, MoE |
| 层 | embedding, encoder, decoder, criterion, optimizer |
| 精度 | fp32, fp16, int8 |
| 模式 | training, inference |
| 兼容性 | Fairseq, Hugging Face, DeepSpeed |
| 解码算法 | beam search, diverse beam search, sampling, CRF |
| 其他 | 梯度通信量化, 自动调优 GEMM 算法 |

下表展示了不同模型当前支持的运行模式和精度。
| 模型 | fp16 训练 | fp16 推理 | int8 训练 | int8 推理 |
| ---------- | --------- | --------- | --------- | --------- |
| Transformer| 是 | 是 | 是 | 是 |
| BERT | 是 | 是 | 是 | 是 |
| GPT2 | 是 | 是 | 是 | 是 |
| BART | 是 | 是 | - | - |
| T5 | - | 是 | - | - |
| MT5 | - | 是 | - | - |
| XGLM | - | 是 | - | - |
| ViT | 是 | 是 | 是 | 是 |
| VAE | - | 是 | - | - |
| 多语言 | - | 是 | - | 是 |
| MoE | - | 是 | - | - |

性能表现

我们在 Transformer 和 BERT 模型上测试了 LightSeq 使用 fp16 和 int8 混合精度进行训练和推理的加速效果。基线为 PyTorch fp16 混合精度。训练实验在单张 A100 GPU 上测试，推理实验在八张 A100 GPU 上测试。

更多性能结果请见此处。

Transformer 训练加速

BERT 训练加速

Transformer 推理加速

BERT 推理加速

安装

通过 PyPI 安装

您可以通过 PyPI 安装 LightSeq，目前仅支持 Linux 系统上的 Python 3.6 至 3.8：

pip install lightseq

从源码构建

您也可以从源码构建：

PATH=/usr/local/hdf5/:$PATH ENABLE_FP32=0 ENABLE_DEBUG=0 pip install -e $PROJECT_DIR

详细的构建说明请见此处。

快速开始

我们在此提供几个示例来展示 LightSeq 的用法。更多详细信息请参阅完整的用户指南和示例。

从零开始使用 LightSeq 训练

您可以使用 LightSeq 提供的模块构建自己的模型。以下是一个构建 Transformer 编码器层的示例。

首先，导入 LightSeq Transformer 编码器模块：

from lightseq.training import LSTransformerEncoderLayer

然后创建一个编码器配置，并用该配置初始化一个 LightSeq Transformer 编码器层：

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=True,
    activation_fn="relu",
    fp16=True,
    local_rank=0,
)
layer = LSTransformerEncoderLayer(config)

除了编码器层，其他模块也可以使用类似的方法创建，然后像普通的 PyTorch 模型一样进行训练。

更多用法请见此处。

基于 Fairseq 使用 LightSeq 训练

LightSeq 将所有快速和轻量的模块集成到了 Fairseq 中。

首先安装以下两个依赖：

pip install fairseq==0.10.2 sacremoses

您可以通过以下命令在 wmt14 en2de 数据集上训练一个 fp16 混合精度的翻译任务：

sh examples/training/fairseq/ls_fairseq_wmt14en2de.sh

（可选）然后您可以在 fp16 预训练模型的基础上开始 int8 混合精度训练：

sh examples/training/fairseq/ls_fairseq_quant_wmt14en2de.sh

更多用法请见此处。

基于 Hugging Face BERT 使用 LightSeq 训练

LightSeq 将 Hugging Face BERT 的编码器层替换为 LightSeq 的快速层。

首先您需要安装这些依赖：

pip install transformers seqeval datasets

在进行下一步训练之前，您需要切换到以下目录：

cd examples/training/huggingface/bert

然后您可以轻松地为不同任务微调 BERT。以命名实体识别任务为例，您可以使用以下命令进行 fp16 混合精度训练：

python task_ner/run_ner.sh

（可选）您也可以在 fp16 预训练模型的基础上开始 int8 混合精度训练：

python task_ner/run_quant_ner.sh

更多用法请见此处。

基于 Fairseq 使用 LightSeq 推理

使用上述脚本训练后，您可以使用 LightSeq 快速进行模型推理。

您需要将 fp16 PyTorch 权重转换为 LightSeq protobuf 或 HDF5 格式：

python export/fairseq/ls_fs_transformer_export.py

（可选）您也可以将 int8 PyTorch 权重转换为 LightSeq protobuf 或 HDF5 格式：

python export/fairseq/ls_fs_quant_transformer_export.py

获得 LightSeq 权重后，您可以使用以下代码快速进行推理：

import lightseq.inference as lsi
model = lsi.Transformer(MODEL_PATH, MAX_BATCH_SIZE)
results = model.infer([[63, 47, 65, 1507, 88, 74, 10, 2057, 362, 9, 284, 6, 2, 1]])

其中 MODEL_PATH 是您的 LightSeq 权重路径，MAX_BATCH_SIZE 是输入句子的最大批次大小。

您也可以通过将 lsi.Transformer 替换为 lsi.QuantTransformer 来快速推理 int8 LightSeq 权重。

更多用法请见此处。

基于 Hugging Face BERT 使用 LightSeq 推理

我们提供了一个端到端的 bert-base 示例，以展示 Lightseq 相比原始 Hugging Face 的速度提升。

首先您需要安装依赖并定位到指定目录：

pip install transformers
cd examples/inference/python

然后您可以通过简单地运行以下命令来检查性能。hf_bert_export.py 用于将 PyTorch 权重转换为 LightSeq protobuf 或 HDF5 格式。

python export/huggingface/hf_bert_export.py
python test/ls_bert.py

更多用法请见此处。

使用推理服务器部署 LightSeq

我们提供了一个包含 tritonserver 和 LightSeq 动态链接库的 Docker 镜像，您只需将模型文件替换为您自己的模型文件即可部署推理服务器。

sudo docker pull hexisyztem/tritonserver_lightseq:22.01-1

更多用法请见此处。

引用我们

如果您在研究中使用了 LightSeq，请引用以下论文。

@InProceedings{wang2021lightseq,
    title = "{L}ight{S}eq: A High Performance Inference Library for Transformers",
    author = "Wang, Xiaohui and Xiong, Ying and Wei, Yang and Wang, Mingxuan and Li, Lei",
    booktitle = "Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies: Industry Papers (NAACL-HLT)",
    month = jun,
    year = "2021",
    publisher = "Association for Computational Linguistics",
    pages = "113--120",
}

@article{wang2021lightseq2,
  title={LightSeq2: Accelerated Training for Transformer-based Models on GPUs},
  author={Wang, Xiaohui and Xiong, Ying and Qian, Xian and Wei, Yang and Li, Lei and Wang, Mingxuan},
  journal={arXiv preprint arXiv:2110.05722},
  year={2021}
}

我们正在招聘！

LightSeq 团队正在招聘具有 深度学习系统、自然语言处理、计算机视觉、语音等 背景的实习生和全职员工。我们位于北京和上海。如果您感兴趣，请将简历发送至 wangxiaohui.neo@bytedance.com。