支援 FlashAttention

支援 FlashAttention 是僅適用分散式轉換器模型程式庫的功能，該模型是以 smp.DistributedModel() 包裝的轉換器模型，用於模型平行訓練。此功能也相容 張量平行處理。

僅當 attention_head_size 所設定的值為 8 的倍數且小於 128 時，FlashAttention 程式庫才會支援模型。因此，當您訓練分散式轉換器並確保 FlashAttention 正常運作時，您應調整參數，以便注意頭大小符合要求。如需更多資訊，另請參閲 FlashAttention GitHub 儲存庫的安裝與功能。

例如，假設您使用 hidden_width=864 與 num_heads=48 設定轉換器模型。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 的倍數)。