Mini-Sequence Transformer: Optimizing Intermediate Memory for Long Sequences Training

要約

非常に長いシーケンスを使用した高効率かつ正確な LLM トレーニングのためのシンプルで効果的な方法論である Mini-Sequence Transformer (MsT) を紹介します。
MsT は入力シーケンスを分割し、ミニシーケンスを繰り返し処理して中間メモリ使用量を削減します。
アクティベーションの再計算と統合されているため、前方パスと後方パスの両方でメモリを大幅に節約できます。
MsT を使用した Llama3-8B モデルの実験では、標準実装よりも 12 倍長いシーケンスでもスループットや収束の低下は測定されませんでした。
MsT は完全に汎用的で実装に依存せず、既存の LLM トレーニング フレームワークと統合するために最小限のコード変更が必要です。
ハグフェイス ライブラリと統合された MsT は、Qwen、Mistral、Gemma-2 の最大コンテキスト長を 12 ~ 24 倍に拡張することに成功しました。

要約(オリジナル)

We introduce Mini-Sequence Transformer (MsT), a simple and effective methodology for highly efficient and accurate LLM training with extremely long sequences. MsT partitions input sequences and iteratively processes mini-sequences to reduce intermediate memory usage. Integrated with activation recomputation, it enables significant memory savings in both forward and backward passes. In experiments with the Llama3-8B model, with MsT, we measure no degradation in throughput or convergence even with 12x longer sequences than standard implementations. MsT is fully general, implementation-agnostic, and requires minimal code changes to integrate with existing LLM training frameworks. Integrated with the huggingface library, MsT successfully extends the maximum context length of Qwen, Mistral, and Gemma-2 by 12-24x.

arxiv情報

著者 Cheng Luo,Jiawei Zhao,Zhuoming Chen,Beidi Chen,Anima Anandkumar
発行日 2024-10-31 16:36:27+00:00
arxivサイト arxiv_id(pdf)

提供元, 利用サービス

arxiv.jp, Google

カテゴリー: cs.LG パーマリンク