[論文レビュー] Efficient Sequence Packing without Cross-contamination: Accelerating Large Language Models without Impacting Performance

Mario Michael Krell, Matej Kosec+2 more|arXiv (Cornell University)|Jun 29, 2021

Natural Language Processing Techniques被引用数 6

ひとこと要約

この論文は、大規模言語モデルの学習におけるパディングトークンを排除するために、シーケンスパッキングをバインパッキング問題として形式化する新しい手法を提案している。パディングを排除することで、性能に損なわれることなく最大2倍の高速化を達成する。効率的で数学的に同等のパッキングアルゴリズムを用い、自己注意機構を変更してシーケンス間の混合を防ぐことで、標準的なハードウェアでも高スループットでの学習が可能になり、精度の低下なしに運用可能である。

ABSTRACT

Effective training of today's large language models (LLMs) depends on large batches and long sequences for throughput and accuracy. To handle variable-length sequences on hardware accelerators, it is common practice to introduce padding tokens, so that all sequences in a batch have the same length. We show in this paper that the variation in sequence lengths in common NLP datasets is such that up to 50% of all tokens can be padding. In less common, but not extreme, cases (e.g. GLUE-cola with sequence length 128), the ratio is up to 89%. Existing methods to address the resulting inefficiency are complicated by the need to avoid cross-contamination in self-attention, by a reduction in accuracy when sequence ordering information is lost, or by customized kernel implementations only valid for specific accelerators. This paper introduces a new formalization of sequence packing in the context of the well-studied bin packing problem, and presents new algorithms based on this formulation which, for example, confer a 2x speedup for phase 2 pre-training in BERT. We show how existing models can be adapted to ensure mathematical equivalence between the original and packed models, meaning that packed models can be trained with existing pre-training and fine-tuning practices.

研究の動機と目的

• 大規模言語モデルの可変長シーケンスバッチ処理におけるパディングトークンに起因する非効率を解消すること。
• 分離トークンやソートバッチ処理に依存する既存手法が抱える、混合の問題、精度の低下、ハードウェア依存の最適化といった制限を克服すること。
• 一般用途で利用可能で、アクセラレータに依存しないシーケンスパッキングフレームワークを構築し、標準的な学習およびファインチューニングワークフローと互換性を持たせること。
• パッキング済みとアンパッキング済みのモデルの間で数学的に同等となるように保証し、収束特性とモデル性能を維持すること。
• 分離トークンや複雑なカーネル実装を必要とせず、近似的に最適なパッキング効率を達成すること。

提案手法

• パディングを最小限に抑えるために、シーケンスパッキングを新しい数学的定式化に基づくバインパッキング問題として形式化する。
• 既存のソルバー（例：非負の最小二乗法）を応用した、決定的で効率的な2つのアルゴリズムを導入し、最適なパッキング戦略を計算する。
• 残差ペナルティを導入した重み付き線形計画法の定式化により、パディングトークンを最小化し、残りのシーケンスを適切に処理する。
• パックされたバッチ内のシーケンス間で混合が生じないよう、自己注意機構を変更し、注意計算をシーケンス内に限定する。
• 残りのシーケンスに対応するため、動的に新しいパッキング戦略を追加する残差ベースの補正ステップを導入する。
• 注意マスクと位置埋め込みを調整することで、パッキング済みシーケンス内の注意パターンを維持し、モデルの同等性を保証する。

実験結果

リサーチクエスチョン

• RQ1 パディングを最小限に抑えつつモデル性能を維持するため、シーケンスパッキングをバインパッキング問題として形式化できるか？
• RQ2 分離トークンやソートバッチ処理に依存する既存のパッキング手法の、パフォーマンスと精度のトレードオフは何か？
• RQ3 ハードウェアに依存しない最適化を必要とせず、標準的な学習パイプラインと互換性を持つパッキング戦略を設計できるか？
• RQ4 パックされたシーケンスにおいて、精度の低下を招かずに自己注意機構における混合を防止する方法は何か？
• RQ5 提案手法は、収束特性とモデル精度を維持したまま、どの程度学習スループットを向上させられるか？

主な発見

• Wikipediaの事前学習データセットでは、シーケンス長512の場合、最大50%のトークンがパディングトークンであり、GLUE-colaでは長さ128で最大89%がパディングである。
• 提案手法により、BERTのフェーズ2事前学習において、パディングトークンを排除することで理論的高速化が最大2倍達成可能である。
• BERT-largeがパッキング済みデータセットで学習する際の収束特性は、アンパッキング済みデータセットと数学的に同等であり、性能の低下はない。
• パディングトークンをほぼゼロにまで削減し、Wikipediaデータセットではパッキング効率が98.7%に達した。
• 数百万のシーケンスを含むデータセットに対しても、1秒未満で最適なパッキング戦略を計算可能であり、リアルタイムの前処理が可能である。
• このアプローチはハードウェアに依存せず、標準的な学習およびファインチューニングワークフローと互換性があり、モデルの再学習やカスタムカーネルの必要がない。