Hugging Faceのトレーニング効率をFlash Attention 2とパッキングで向上させる
Improving Hugging Face Training Efficiency Through Packing with Flash Attention 2
Hugging Face Transformersは、パディングなしでインストラクションチューニングの例をパッキングする新機能とDataCollatorWithFlatteningを導入し、Flash Attention 2との互換性を実現しました。これにより、トレーニングのスループットが最大2倍向上し、ピークメモリ使用量も削減されますが、収束の質は維持されます。この機能は、シーケンス長に大きなばらつきがあるデータセット(例: FLAN)で特に効果的で、スループットが2倍になることがあります。一方、シーケンス長が比較的均一なデータセット(例: OrcaMath)では1.4倍のスループット向上が見られます。メモリ使用量はFLANで20%、OrcaMathで6%削減されます。この機能はposition_idsを公開する14のモデル(Llama 2/3、Mistral、Gemmaなど)で利用可能です。TRLライブラリのSFTTrainerユーザーは、DataCollatorForCompletionOnlyLMでpadding_free=Trueを設定することで同様の恩恵を受けられます。
- Hugging Face Transformersがパディングなしでのインストラクションチューニング用パッキング機能とDataCollatorWithFlatteningを導入し、Flash Attention 2との互換性を実現した
- 本機能によりトレーニングスループットが最大2倍向上し、ピークメモリ使用量が最大20%削減される
Hugging FaceがFlash Attention 2対応とパッキング機能によりトレーニング効率を最大2倍向上させたことは、LLMのファインチューニングコスト削減に直結する。特にデータセットのシーケンス長にばらつきがある実用的なシナリオで効果が大きく、GPUリソースの逼迫するスタートアップにとってはトレーニング時間短縮・メモリ削減という明確な経済的便益がある。Hugging Faceエコシステムの競争力強化にも寄与し、同社のプラットフォーム価値向上につながる。