Support for FlashAttention

Support for FlashAttention é um recurso da biblioteca aplicável apenas ao modelo de transformador distribuído, que é um modelo de transformador incluído smp.DistributedModel()para treinamento paralelo de modelos. Esse recurso também é compatível com Paralelismo tensorial.

A FlashAttentionbiblioteca só oferece suporte a modelos quando attention_head_size é definida com um valor múltiplo de 8 e menor que 128. Portanto, ao treinar um transformador distribuído e garantir que ele FlashAttention funcione corretamente, você deve ajustar os parâmetros para que o tamanho da cabeça de atenção atenda aos requisitos. Para obter mais informações, consulte também Instalação e recursos no FlashAttention GitHubrepositório.

Por exemplo, suponha que você configure um modelo Transformador com hidden_width=864 e num_heads=48. O tamanho da cabeça de FlashAttention é calculado comoattention_head_size = hidden_width / num_heads = 864 / 48 = 18. Para habilitar FlashAttention, você precisa ajustar o num_heads parâmetro para54, de forma que attention_head_size = hidden_width / num_heads = 864 / 54 = 16 seja um múltiplo de 8.