FlashAttention のサポート

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

FlashAttentionattention_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パラメータを 8 attention_head_size = hidden_width / num_heads = 864 / 54 = 16 の倍数になるように調整する必要があります。54