Résumé Technique : AdaptiveLoad

1. Énoncé du Problème

L'entraînement de Transformers de diffusion vidéo à grande échelle (tels que DiT et MMDiT) pour les modèles du monde se heurte à deux goulots d'étranglement critiques lors du traitement de jeux de données en mode mixte contenant des longueurs de séquence hautement hétérogènes (allant d'images statiques à des vidéos de longue durée) :

• Déséquilibre computationnel dû à la complexité quadratique : Les pipelines d'entraînement distribués existants reposent sur des stratégies de « mise en seau égale en tokens » (contrainte $B \times S = \text{constante}$). Cette approche ne prend pas en compte la complexité computationnelle $O(S^2)$ des mécanismes d'attention auto-attention. Par conséquent, les seaux contenant des séquences longues deviennent des « traînards computationnels », provoquant un déséquilibre de charge sévère. Le temps de synchronisation global est dicté par le GPU le plus lent, entraînant un temps d'attente important (bulles de synchronisation) pour les autres nœuds. Les données expérimentales indiquent une faible corrélation ($R \approx 0,35$) entre la latence de l'étape et le nombre total de tokens, mais une forte corrélation ($R \approx 0,92$) avec le carré de la longueur de la séquence.
• Mur de mémoire et surcharge des opérateurs : Dans les architectures Multimodal DiT (MMDiT), la Normalisation de Couche Adaptative (AdaLN) est invoquée des centaines de fois par itération. Les implémentations standard utilisent des noyaux CUDA discrets, provoquant des allers-retours fréquents et redondants des tenseurs intermédiaires entre les Multiprocesseurs de Flux (SM) et la Mémoire à Large Bande Passante (HBM). De plus, lors de la rétropropagation, la réduction des gradients pour les paramètres de modulation souffre de motifs d'accès mémoire à pas (accès non contigus le long de la dimension de séquence $N$), empêchant l'accès coalescé au niveau matériel et sous-utilisant sévèrement la bande passante mémoire.

2. Méthodologie : Le Cadre AdaptiveLoad

L'article propose AdaptiveLoad, un cadre d'optimisation full-stack intégrant une planification de haut niveau avec la fusion de noyaux de bas niveau.

A. Équilibrage de Charge Adaptatif à Double Contrainte

Pour remédier au déséquilibre computationnel, les auteurs remplacent l'hypothèse linéaire « égale en tokens » par une logique de planification à double contrainte qui impose simultanément :

1. Contrainte Linéaire de Mémoire : $B \times S \leq M_{mem}$ (assurant que la taille du lot rentre dans la mémoire du GPU).
2. Contrainte Computationnelle Polynomiale : $B \times S^p \leq M_{comp}$ (tenant compte de la complexité $O(S^2)$).

La taille finale du lot pour un seau est déterminée par l'intersection de ces contraintes :
$$B_{shape} = \max(1, \min(\lfloor M_{mem}/S \rfloor, \lfloor M_{comp}/S^p \rfloor))$$

• Système d'ajustement des coûts : Pour déterminer l'exposant empirique $p$ et la borne computationnelle $M_{comp}$, le système emploie un flux de travail de Référencement de Forme (Shape Benchmarking) piloté par les données. Il exécute des scans de pixels synthétiques dans un environnement distribué en direct (par exemple, les chemins FSDP) pour mesurer la latence de l'étape en fonction de la taille du lot et de la longueur de la séquence. En utilisant une régression non linéaire, il ajuste un modèle de coûts ($\text{temps d'étape} \approx a + b \times B \times S^p$) pour calibrer dynamiquement les configurations des seaux, éliminant ainsi le besoin d'un réglage empirique manuel.

B. Noyaux CUDA Fusionnés LayerNorm-Modulate

Pour remédier au « Mur de mémoire », les auteurs conçoivent un noyau CUDA AdaLN fusionné qui combine les opérations de normalisation et de modulation en un seul lancement de noyau.

• Passage Avant : Le noyau utilise les registres et la Mémoire Partagée pour stocker les statistiques intermédiaires (moyenne et variance), effectuant des réductions en deux étapes (mélange au niveau du Warp et réduction au niveau du Bloc). Les transformations linéaires sont appliquées directement lors d'un troisième parcours de la dimension de caractéristiques $D$, éliminant les lectures/écritures HBM redondantes.
• Passage Arrière (Réduction Coalescée D-tile) : Pour résoudre le goulot d'étranglement de l'accès à pas dans la réduction des gradients, les auteurs introduisent une stratégie de Réduction Coalescée D-tile. Au lieu de réduire le long de la dimension de séquence $N$ (ce qui provoque un accès à pas), la Grille est partitionnée en une disposition $(d_{tile}, n_{tile})$. Les threads sont fixés à un indice de caractéristique $d$ et effectuent une accumulation verticale le long des tuiles de séquence. Cela garantit que les accès mémoire au sein d'un Warp sont contigus, permettant des transactions HBM coalescées et maximisant l'utilisation de la bande passante.
• Simplification du Graphique : En fusionnant la chaîne d'opérateurs, le graphique de calcul est effondré de 5 à 8 nœuds discrets en un seul nœud atomique. Les statistiques intermédiaires restent dans les registres, contournant la gestion de la mémoire globale du cadre et réduisant considérablement l'empreinte mémoire des activations.

3. Contributions Clés

1. Quantification et Atténuation du Déséquilibre de Charge : L'article quantifie rigoureusement le déséquilibre de charge dans l'entraînement de la diffusion vidéo et propose un algorithme de planification de seaux adaptatif à double contrainte. Cette approche réduit le taux de déséquilibre de charge jusqu'à 50 % en gérant activement la complexité $O(S^2)$ des séquences longues.
2. Modélisation des Coûts Pilotée par les Données : Les auteurs conçoivent un système de bout en bout qui fait passer la configuration de l'entraînement distribué d'une approche heuristique à une approche pilotée par les données. Grâce à un ajustement automatique du modèle de coûts via le Référencement de Forme, le système optimise l'utilisation des ressources à l'échelle du cluster sans intervention manuelle.
3. Noyaux Fusionnés Haute Performance : Un noyau CUDA AdaLN fusionné spécialisé est développé pour les architectures de base MMDiT. En introduisant une réduction coalescée D-tile, il résout le goulot d'étranglement de l'accès à pas dans la rétropropagation, augmentant le débit d'entraînement de plus de 25 % tout en réduisant simultanément l'utilisation de la mémoire d'activation.

4. Résultats Expérimentaux

Les expériences ont été menées sur le modèle du monde Wan 2.1 utilisant un corpus mixte de 10 millions d'échantillons (WebDataset et Koala-36m) sur des clusters de 8 et 16 GPU.

• Amélioration du Débit :
  • Cluster de 8 GPU : Le débit moyen d'entraînement est passé de 14 383 à 18 069 tokens/seconde (+25,6 %).
  • Cluster de 16 GPU : Le débit moyen est passé de 30 170 à 38 372 tokens/seconde (+27,18 %).
  • La courbe AdaptiveLoad a maintenu un plancher de débit plus élevé par rapport à la référence, même dans les distributions de séquences pires cas.
• Efficacité de l'Équilibrage de Charge :
  • Le taux de déséquilibre computationnel (mesuré par le coefficient de variation de la charge de calcul) a été réduit de 39 % à 18,9 %.
  • La surcharge de synchronisation (CV du temps d'étape) a chuté de manière significative, la configuration 16 GPU passant d'une moyenne de référence de 18,7 % à 10,4 % sous AdaptiveLoad.
• Performance Mémoire et Noyau :
  • Utilisation VRAM Pic : Améliorée de 22,7 %.
  • Accélération du Noyau : Le passage avant fusionné a atteint une accélération de 3,21× à 3,39×. Le passage arrière (réduction D-tile) a atteint une accélération allant jusqu'à 1,42× pour une longueur de séquence de 64k.
  • Économies de Mémoire : Le noyau fusionné a réduit la mémoire d'activation d'environ 61,9 % (par exemple, économisant ~4,8 Go pour $N=64 000$).
  • Capacité de Charge : Le système optimisé a pris en charge une longueur de séquence maximale de 52 800 tokens (une augmentation de 10 % par rapport à la limite de 48 000 tokens de la référence) sans augmenter l'utilisation de la mémoire pic.
• Stabilité de l'Entraînement : Les courbes de perte pour AdaptiveLoad sont restées hautement congruentes avec la référence, démontrant que les optimisations ont préservé la fidélité numérique et la qualité de la convergence.

5. Importance

L'article affirme qu'AdaptiveLoad fournit une base robuste pour la mise à l'échelle des modèles de génération vidéo de séquences ultra-longues. En conciliant la sécurité de la mémoire avec la parité computationnelle grâce à une planification à double contrainte et en s'attaquant au « Mur de mémoire » via une fusion d'opérateurs consciente du matériel, le cadre élimine efficacement les goulots d'étranglement de synchronisation inhérents aux jeux de données vidéo hétérogènes. Les résultats démontrent que dépasser les heuristiques de comptage linéaire de tokens pour une planification consciente de la complexité quadratique est essentiel pour maximiser l'efficacité des clusters GPU à grande échelle dans les tâches de génération vidéo.