Support FlashAttention - Amazon SageMaker
Amazon Web Services 文档中描述的 Amazon Web Services 服务或功能可能因区域而异。要查看适用于中国区域的差异,请参阅 中国的 Amazon Web Services 服务入门 (PDF)

本文属于机器翻译版本。若本译文内容与英语原文存在差异,则一律以英文原文为准。

Support FlashAttention

Suppor FlashAttention t for 是该库的一项功能,仅适用于分布式变压器模型,分布式变压器模型是为模型并行训练而封装的 Trans smp.DistributedModel()former 模型。此功能还与 张量并行性 兼容。

FlashAttention库仅在设置attention_head_size为 8 的倍数且小于 128 的值时才支持模型。因此,在训练分布式变压器并确保其 FlashAttention 正常工作时,应调整参数以使注意力头大小符合要求。有关更多信息,另请参阅FlashAttention GitHub存储库中的安装和功能

例如,假设您使用 hidden_width=864num_heads=48 配置转换器模型。的头部大小计算公式 FlashAttention 为attention_head_size = hidden_width / num_heads = 864 / 48 = 18。要启用 FlashAttention,您需要将num_heads参数调整为 54attention_head_size = hidden_width / num_heads = 864 / 54 = 16,即 8 的倍数。