Increasing the context length of large language models (LLMs) unlocks
fundamentally new capabilities, but also significantly increases the memory
footprints of training. Previous model-parallel systems such as Megatron-LM
partition and compute different attention heads in parallel, resulting in large
communication volumes, so they cannot scale beyond the number of attention
heads, thereby hindering its adoption. In this paper, we introduce a new
approach, LightSeq, for long-context LLMs training. LightSeq has many notable
advantages. First, LightSeq partitions over the sequence dimension, hence is
agnostic to model architectures and readily applicable for models with varying
numbers of attention heads, such as Multi-Head, Multi-Query and Grouped-Query
attention. Second, LightSeq not only requires up to 4.7x less communication
than Megatron-LM on popular LLMs but also overlaps the communication with
computation. To further reduce the training time, LightSeq features a novel
gradient checkpointing scheme to bypass an forward computation for
memory-efficient attention. We evaluate LightSeq on Llama-7B and its variants
with sequence lengths from 32K to 512K. Through comprehensive experiments on
single and cross-node training, we show that LightSeq achieves up to 1.24-2.01x
end-to-end speedup, and a 2-8x longer sequence length on models with fewer
heads, compared to Megatron-LM. Codes will be available at
this https URL

LightSeq 是一种新的方法，用于长上下文大语言模型 (LLMs) 的训练，在流行的 LLMs 上比 Megatron-LM 通信量少且重叠计算，通过新的梯度检查点方案实现高效的注意力计算。

LightSeq：面向长上下文 Transformer 的分布式训练的序列级并行

LightSeq: Sequence Level Parallelism for Distributed Training of Long  Context Transformers

Transformer, BERT and their variants have achieved great success in natural
language processing. Since Transformer models are huge in size, serving these
models is a challenge for real industrial applications. In this paper, we
propose LightSeq, a highly efficient inference library for models in the
Transformer family. LightSeq includes a series of GPU optimization techniques
to to streamline the computation of neural layers and to reduce memory
footprint. LightSeq can easily import models trained using PyTorch and
Tensorflow. Experimental results on machine translation benchmarks show that
LightSeq achieves up to 14x speedup compared with TensorFlow and 1.4x compared
with FasterTransformer, a concurrent CUDA implementation. The code is available
at this https URL

本文提出 LightSeq，一种高效的推断库，以加速 Transformer 系列模型的计算和减少内存占用。实验结果表明，相较于 TensorFlow 和 FasterTransformer，LightSeq 可以实现多达 14 倍的加速。