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 的倍数。