Support per FlashAttention

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

La FlashAttentionlibreria supporta i modelli solo quando attention_head_size è impostata su un valore che è un multiplo di 8 e inferiore a 128. Pertanto, quando si addestra un trasformatore distribuito e ci si assicura che FlashAttention funzioni correttamente, è necessario regolare i parametri per fare in modo che la dimensione della testina di attenzione soddisfi i requisiti. Per ulteriori informazioni, consulta anche Installazione e funzionalità nel FlashAttention GitHubrepository.

Ad esempio, supponiamo di configurare un modello Transformer con hidden_width=864 e num_heads=48. La dimensione della testa di FlashAttention è calcolata comeattention_head_size = hidden_width / num_heads = 864 / 48 = 18. Per abilitarlo FlashAttention, è necessario regolare il num_heads parametro su54, in modo che attention_head_size = hidden_width / num_heads = 864 / 54 = 16 sia un multiplo di 8.