FlashMask-3 accélère l'attention structurée sur plusieurs GPU

Lorsque l'entraînement utilise de longues séquences réparties sur plusieurs GPU, peut-on appliquer des masques d'attention structurée sans que leur traitement et les déséquilibres de charge ne deviennent un nouveau goulot d'étranglement ?

La question est importante car distribuer la séquence ajoute à la fois du calcul et de la communication au travail que les masques conditionnent.

FlashMask-3 combine un noyau d'attention Hopper avec des masques locaux par rang, un équilibrage du travail épars et une communication chevauchée pour l'entraînement distribué.

FlashMask-3 combine un noyau d'attention pour GPU Hopper avec des méthodes distribuées de localisation des masques, d'équilibrage du travail épars et de chevauchement du calcul avec la communication.

La comparaison a mesuré l'entraînement complet face à l'équilibrage Dual Chunk et à la communication all-gather pour un modèle 118B-A14B à un contexte de 128 000 tokens sur 128 GPU H100.

Le débit d'entraînement était supérieur sur les trois jeux de données à masques de documents, de 7,3 % à 11,0 % selon les jeux de données présentant une sparsité moyenne des masques de 70 % à 86 %.

Ce résultat atteste d'un gain pour la méthode combinée dans cette configuration ; les seuls benchmarks d'attention n'établiraient pas à eux seuls une amélioration du débit d'entraînement complet.

Le gain sur l'entraînement complet a été mesuré pour un seul modèle 118B-A14B, des masques de documents, trois jeux de données et une configuration spécifique à 128 GPU H100.

L'article rapporte également que la désactivation du routage hiérarchique a fait passer le gain de la méthode combinée en dessous de la baseline sur deux de ces jeux de données.