Supporto per FlashAttention

Il supporto per FlashAttention è una funzionalità della libreria applicabile solo al modello di trasformatore distribuito, che è un modello Transformer integrato da smp.DistributedModel() per l'addestramento parallelo al modello. Questa funzionalità è compatibile anche con Parallelismo tensoriale.

La libreria FlashAttention supporta i modelli solo se attention_head_size è impostato su un valore che è un multiplo di 8 e inferiore a 128. Pertanto, quando addestri un trasformatore distribuito e ti assicuri che FlashAttention funzioni correttamente, dovresti regolare i parametri per far sì che la dimensione delle “head attention” soddisfi i requisiti. Per ulteriori informazioni, consulta anche Installazione e funzionalità nel repository GitHub di FlashAttention.

Ad esempio, supponiamo di configurare un modello Transformer con hidden_width=864 e num_heads=48. La dimensione della “head attention” di FlashAttention viene calcolata come attention_head_size = hidden_width / num_heads = 864 / 48 = 18 Per abilitare FlashAttention, è necessario regolare il parametro num_headssu 54, in modo che attention_head_size = hidden_width / num_heads = 864 / 54 = 16 sia un multiplo di 8.