DISTFLASHATTN: Distributed Memory-efficient Attention for Long-context LLMs Training
下記の文章は論文のアブストラクトをChatGPTに解釈させたものであり、正確性は保証しません。
本研究は、長い文脈を持つ大規模な言語モデル（LLMs）のトレーニングにおいて、メモリ使用量を効率的に削減するための分散メモリ効率の高いアテンションメカニズムであるDISTFLASHATTNを紹介しています。先行研究では、単一のGPU上でのアテンション操作のピークメモリ使用量を削減するためのメモリ効率の高いアテンションが提案されていますが、分散拡張がないため、単一のデバイスで処理できるシーケンス長に限定されています。本研究では、FlashAttentionの利点を保ちながら、分散設定に拡張するためのDISTFLASHATTNを提案しています。
本研究の目的は、長い文脈を持つLLMsのトレーニングにおいて、メモリ効率の高いアテンションを分散設定で利用することです。具体的には、トークンレベルのワークロードのバランス調整、キーと値の通信のオーバーラップ、および再素材化に対応した勾配チェックポイントアルゴリズムの3つの技術を提案しています。これらの技術により、DISTFLASHATTNは長いシーケンスに対して8倍の長さ、Ring Self-Attentionと比較して4.45〜5.64倍の高速化、Megatron-LM with FlashAttentionと比較して2〜8倍の長さ、1.24〜2.01倍の高速化を実現しています。
本研究で使用したデータについては、具体的な情報は提供されていません。
本研究で使用した手法は、DISTFLASHATTNと呼ばれる分散メモリ効率の高いアテンションメカニズムです。この手法では、トークンレベルのワークロードのバランス調整、キーと値の通信のオーバーラップ、および再素材化に対応した勾配チェックポイントアルゴリズムの3つの技術が使用されています。これにより、DISTFLASHATTNは長いシーケンスに対して効率的に処理を行うことができます。
本研究によって明らかになったことは、DISTFLASHATTNが長いシーケンスに対して高い効率性を持ち、他の分散システムと比較して8倍の長さ、1.67倍の高速化を実現できることです。また、DISTFLASHATTNはRing AttentionやDeepSpeed-Ulyssesと比較しても1.26〜1.88倍の高速化を実現しています。