Support para FlashAttention

Support for FlashAttention es una función de la biblioteca que solo se aplica al modelo de transformador distribuido, que es un modelo de transformador incluido smp.DistributedModel()para el entrenamiento en paralelo con modelos. Esta característica también es compatible con Paralelismo de tensores.

La FlashAttentionbiblioteca solo admite modelos cuando attention_head_size se establece en un valor que es múltiplo de 8 e inferior a 128. Por lo tanto, cuando entrenes un transformador distribuido y te asegures de que FlashAttention funciona correctamente, debes ajustar los parámetros para que el tamaño del cabezal de atención cumpla con los requisitos. Para obtener más información, consulte también Instalación y características del FlashAttention GitHubrepositorio.

Por ejemplo, supongamos que configura un modelo de Transformer con hidden_width=864 y num_heads=48. El tamaño de la cabeza de FlashAttention se calcula comoattention_head_size = hidden_width / num_heads = 864 / 48 = 18. Para FlashAttention activarlo, debe ajustar el num_heads parámetro a54, de modo que attention_head_size = hidden_width / num_heads = 864 / 54 = 16 sea un múltiplo de 8.