Linear Attention Sequence Parallelism

要約

シーケンス並列(SP)は、単一のGPUのメモリ限界を超える長いシーケンスを処理するための一般的な戦略として機能する。しかし、既存のSP手法は線形注意の特徴を利用していないため、線形注意ベースの言語モデルにとって最適な並列効率と使い勝手が得られていません。本論文では、線形注意ベースの言語モデルに合わせた効率的なSP手法である線形注意シーケンス並列(LASP)を紹介する。具体的には、線形注意の右積カーネルトリックを活用する効率的なポイントツーポイント通信機構を設計し、SPの通信オーバヘッドを大幅に削減する。また、カーネルフュージョンと中間状態のキャッシュを行うことで、LASPの実用的な効率を高め、GPUクラスタ上でのLASPの実装をハードウェアフレンドリーにする。さらに、シーケンスレベルのLASPとバッチレベルのデータ並列手法の互換性を綿密に保証する。我々は、シーケンス長とGPUクラスタサイズを変化させた2つの線形注意ベースのモデルで広範な実験を行った。LASPは、1Bモデルで128個のA100 80G GPUを使い、シーケンス長を4096Kまで拡張し、既存のSP法より8倍長く、かつ大幅に高速化した。コードはhttps://github.com/OpenNLPLab/LASP。

要約(オリジナル)

Sequence Parallel (SP) serves as a prevalent strategy to handle long sequences that exceed the memory limit of a single GPU. However, existing SP methods do not take advantage of linear attention features, resulting in sub-optimal parallelism efficiency and usability for linear attention-based language models. In this paper, we introduce Linear Attention Sequence Parallel (LASP), an efficient SP method tailored to linear attention-based language models. Specifically, we design an efficient point-to-point communication mechanism to leverage the right-product kernel trick of linear attention, which sharply decreases the communication overhead of SP. We also enhance the practical efficiency of LASP by performing kernel fusion and intermediate state caching, making the implementation of LASP hardware-friendly on GPU clusters. Furthermore, we meticulously ensure the compatibility of sequence-level LASP with all types of batch-level data parallel methods, which is vital for distributed training on large clusters with long sequences and large batches. We conduct extensive experiments on two linear attention-based models with varying sequence lengths and GPU cluster sizes. LASP scales sequence length up to 4096K using 128 A100 80G GPUs on 1B models, which is 8 times longer than existing SP methods while being significantly faster. The code is available at https://github.com/OpenNLPLab/LASP.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL