FlashAttention のサポート

FlashAttention のサポートはライブラリの機能で、モデル並列トレーニングの smp.DistributedModel() によってラップされた Transformer モデルである分散トランスフォーマーモデルにのみ適用されます。この機能は テンソル並列処理 とも互換性があります。

FlashAttention ライブラリは、attention_head_size が 8 の倍数で 128 未満の値に設定されているモデルのみをサポートしています。したがって、分散トランスフォーマーをトレーニングして FlashAttention が正しく動作することを確認するときは、アテンションヘッドサイズが要件を満たすようにパラメータを調整する必要があります。詳細については、「FlashAttention GitHub リポジトリ」の「インストールと機能」も参照してください。

例えば、hidden_width=864 と num_heads=48 を使用して Transformer モデルを設定すると仮定します。FlashAttention のヘッドサイズは attention_head_size = hidden_width / num_heads = 864 / 48 = 18 と計算されます。FlashAttention を有効にするには、num_heads パラメータを 54 に調整して attention_head_size = hidden_width / num_heads = 864 / 54 = 16 (つまり 8 の倍数) となるように調整する必要があります。