💬 NLP

Folding Tensor and Sequence Parallelism for Memory-Efficient Transformer Training & Inference

Ce papier présente la parallélisation tensorielle et séquentielle (TSP), une nouvelle stratégie d'exécution qui combine le fractionnement des poids et des jetons sur un seul axe de dispositif pour réduire simultanément la surcharge mémoire des paramètres et des activations, offrant ainsi une alternative efficace sur le plan matériel pour l'entraînement et l'inférence de modèles de transformateurs à contexte long et à contraintes mémoire.

Auteurs originaux : Vasu Shyam, Anna Golubeva, Quentin Anthony

Publié 2026-04-30

📖 6 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Vasu Shyam, Anna Golubeva, Quentin Anthony

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous essayez de résoudre un immense puzzle avec un groupe d'amis, mais que vous ne disposez que d'une très petite table (la mémoire de votre ordinateur) pour travailler. Le puzzle est si grand qu'aucune personne seule ne peut tenir toutes les pièces à la fois.

Ce document présente une nouvelle méthode permettant à une équipe d'ordinateurs (GPU) de travailler ensemble pour entraîner des modèles d'IA gigantesques, qui sont essentiellement ces immenses puzzles. Les auteurs appellent leur nouvelle stratégie TSP (Parallélisme Tensoriel et Séquentiel).

Voici le détail utilisant des analogies simples :

Le Problème : Deux Vieilles Façons de Partager le Travail

Pour résoudre le puzzle, l'équipe utilise généralement l'une des deux vieilles méthodes, mais toutes deux présentent des défauts :

La Méthode « Division des Poids » (Parallélisme Tensoriel) :
Imaginez que les pièces du puzzle sont les « règles » du jeu (les poids du modèle). Dans cette méthode, vous coupez le livre de règles en deux. La personne A tient la première moitié des règles, et la personne B tient la seconde moitié.
- Le Bon : Vous économisez de la place sur la table car vous ne stockez pas le livre de règles entier deux fois.
- Le Mauvais : Si le puzzle contient une longue histoire (une longue séquence de mots), tout le monde doit encore tenir l'histoire entière dans ses mains pour jouer. Si l'histoire est énorme, vos mains (mémoire) se remplissent et vous plantez.
La Méthode « Division de l'Histoire » (Parallélisme Séquentiel) :
Imaginez que les pièces du puzzle sont l'histoire elle-même. Dans cette méthode, la personne A tient la première moitié de l'histoire, et la personne B tient la seconde moitié.
- Le Bon : Vous économisez de la place sur la table car vous ne tenez pas l'histoire entière d'un coup.
- Le Mauvais : Tout le monde doit encore mémoriser le livre de règles entier. Si le livre de règles est énorme, votre cerveau (mémoire) se remplit et vous plantez.

L'Ancien Hybride : Habituellement, les équipes tentent de faire les deux en utilisant deux groupes distincts d'amis. Un groupe divise les règles, et un groupe différent divise l'histoire. Mais cela est inefficace car cela utilise tous vos amis juste pour diviser le travail, ne laissant personne pour aider avec d'autres tâches (comme le Parallélisme des Données).

La Solution : La Méthode « Pliée » (TSP)

Les auteurs disent : « Pourquoi utiliser deux groupes séparés ? Pliions le travail sur un seul axe. »

Dans le TSP, chaque personne du groupe fait les deux choses en même temps :

Elle tient une tranche du livre de règles (poids).
Elle tient une tranche de l'histoire (séquence).

L'Analogie :
Imaginez que vous êtes à un dîner.

Vieille Façon : Vous avez une table où une personne fait circuler le menu (poids) pendant que tout le monde lit le livre entier. Une autre table a des gens qui font circuler le livre (histoire) pendant que tout le monde mémorise le menu entier.
Façon TSP : Tout le monde à la table reçoit un petit morceau du menu et un petit morceau de l'histoire.

Comment Ils Font Fonctionner Cela (Les Tours de Magie)

Puisque chacun a un tout petit morceau du menu et un tout petit morceau de l'histoire, ils doivent beaucoup se parler pour finir le puzzle. Le document décrit deux façons astucieuses de le faire sans être submergés :

Pour les parties « Histoire » (Attention) :
Imaginez que le groupe a besoin de connaître l'histoire entière pour comprendre une phrase spécifique. Au lieu que tout le monde crie l'histoire entière d'un coup, ils se relaient. Une personne diffuse sa partie du menu à tout le monde. Ensuite, tout le monde calcule sa partie de l'histoire, et ils échangent rapidement leurs morceaux d'histoire (clés et valeurs) pour reconstruire le contexte complet. C'est comme une course de relais où ils passent le témoin (données) tout en courant.
Pour les parties « Règles » (MLP) :
Imaginez que le groupe doit appliquer différentes règles à leurs morceaux d'histoire. Au lieu de s'arrêter pour crier les règles, ils font passer les pages du livre de règles en cercle (un anneau). La personne A fait ses calculs avec la page 1, puis passe la page 1 à la personne B pendant que la personne B passe la page 2 à la personne C. Pendant que les pages bougent, tout le monde est occupé à faire des calculs. Cela maintient le « trafic » en mouvement pendant que le « travail » est accompli.

Pourquoi Est-ce Mieux ?

Le document affirme que le TSP est une solution « consciente du matériel », ce qui signifie qu'elle est conçue spécifiquement pour la façon dont les puces informatiques modernes communiquent entre elles.

Économies de Mémoire : Parce que chacun tient un morceau des règles et un morceau de l'histoire, la mémoire requise sur chaque ordinateur diminue considérablement. Cela permet à l'équipe de gérer des histoires beaucoup plus longues (contexte plus long) sans épuiser la mémoire.
Vitesse : Même s'ils échangent plus de données d'avant en arrière (ce qui semble plus lent), ils le font d'une manière qui se superpose à leur réflexion. Le « passage » a lieu pendant qu'ils « réfléchissent », donc le temps total ne s'allonge pas beaucoup.
Tenir dans la Pièce : Dans un cluster informatique, la connexion la plus rapide est généralement entre les puces sur la même machine (comme des gens assis à la même table). La connexion plus lente est entre les machines différentes (des gens dans des pièces différentes).
- Les anciennes méthodes forçaient souvent l'équipe à se diviser entre différentes pièces, les ralentissant.
- Le TSP permet à toute l'équipe « divisée» de tenir sur une seule machine (une seule table), les maintenant dans la voie rapide.

Les Résultats

Les auteurs ont testé cela sur un cluster massif de 1 024 GPU puissants (MI300X).

Mémoire : Le TSP a utilisé la moindre quantité de mémoire à chaque test, surtout lorsque les histoires étaient très longues.
Vitesse : Le TSP était aussi rapide, voire plus rapide, que les anciennes méthodes.
Évolutivité : Alors qu'ils ajoutaient plus d'ordinateurs au groupe, le TSP continuait à bien performer, tandis que les anciennes méthodes commençaient à lutter avec les limites de mémoire.

En bref : Le TSP est une façon plus intelligente d'organiser une équipe d'ordinateurs. Au lieu de diviser les « règles » et l'« histoire » en groupes séparés, il les combine afin que chaque ordinateur détienne un peu des deux. Cela économise de la place, permet des histoires plus longues et maintient l'équipe travaillant efficacement sur le même réseau rapide.

1. Énoncé du problème

L'entraînement et l'inférence des grands modèles de langage (LLM) se heurtent à deux goulots d'étranglement mémoire principaux :

Mémoire des paramètres : Le stockage requis pour les poids du modèle, les gradients et les états de l'optimiseur.
Mémoire des activations : Le stockage requis pour les activations intermédiaires, qui évolue de manière quadratique avec la longueur de séquence ( $S^2$ ) dans les mécanismes d'attention.

Les stratégies de parallélisme existantes abordent ces problèmes de manière isolée ou avec des compromis :

Parallélisme tensoriel (TP) : Il fragmente les poids du modèle sur plusieurs dispositifs, réduisant la mémoire des paramètres mais laissant la mémoire des activations entièrement répliquée. Il devient inefficace pour les longues séquences en raison de la mémoire d'activation élevée et du volume de communication.
Parallélisme de séquence (SP) : Il fragmente la séquence d'entrée sur plusieurs dispositifs, réduisant la mémoire des activations mais laissant les poids du modèle entièrement répliqués. Il est inefficace lorsque la mémoire des paramètres constitue le goulot d'étranglement.
TP+SP conventionnel : Il combine les deux en les assignant à des axes orthogonaux dans une grille de dispositifs (par exemple, une grille $T \times \Sigma$ ). Bien qu'efficace, cela consomme un grand nombre de dispositifs pour le parallélisme de modèle ( $T \times \Sigma$ ), laissant moins de dispositifs pour le parallélisme de données. De plus, cela force souvent les groupes parallèles à s'étendre sur des liens inter-nœuds plus lents, sous-utilisant les interconnexions intra-nœuds à haute bande passante (comme NVLink ou AMD Infinity Fabric).

Le défi central : Comment réduire simultanément la mémoire des paramètres et celle des activations sur un seul axe de dispositif pour maximiser l'efficacité mémoire et libérer des dispositifs pour le parallélisme de données, sans engendrer de coûts de communication prohibitifs.

2. Méthodologie : Parallélisme tensoriel et de séquence (TSP)

Les auteurs proposent le TSP, une stratégie qui « plie » le parallélisme tensoriel et le parallélisme de séquence sur un seul axe de dispositif. Au lieu d'assigner le TP et le SP à des dimensions orthogonales, le TSP assigne à chaque rang (GPU) à la fois une fraction des poids du modèle et une fraction de la séquence.

Composantes clés de la conception :

Fragmentation unifiée : Chaque dispositif détient $1/D$ des paramètres du modèle et $1/D$ des jetons de séquence. Cela réduit à la fois la mémoire proportionnelle aux paramètres et la mémoire des activations d'un facteur $D$ sur un seul axe.
Ordonnancements d'exécution : Pour gérer le fait que les poids et les activations sont fragmentés différemment, le TSP emploie des schémas de communication spécifiques pour différentes couches :
- Couches d'attention : Utilise un schéma de diffusion (broadcast).
  - Les rangs itèrent sur les fragments de poids (empaquetés $W_Q, W_K, W_V, W_O$ ).
  - À chaque itération, un rang diffuse son fragment de poids à tous les autres.
  - Tous les rangs calculent des projections locales sur leurs fragments de séquence.
  - Une opération All-Gather reconstruit le contexte complet Clé/Valeur sur la dimension de séquence pour effectuer FlashAttention.
  - Un schéma de partitionnement en zigzag est utilisé pour équilibrer la charge de travail de l'attention causale (assurant que les jetons ultérieurs n'encourent pas un travail disproportionnellement plus important).
- Couches MLP à porte : Utilise un schéma de communication en anneau.
  - Les fragments de poids ( $W_1, W_2, W_3$ ) circulent de point à point autour de l'anneau.
  - Chaque rang accumule localement les sorties partielles.
  - Cela élimine le besoin d'une opération All-Reduce (standard dans le TP), car la séquence reste locale et seuls les poids se déplacent.
Conscience matérielle : La conception est optimisée pour des topologies comme celle du MI300X d'AMD (Infinity Fabric), où la bande passante collective évolue avec le nombre de GPU participants. Le TSP garantit que les primitives de communication (diffusion et anneau) engagent le groupe complet ou sont superposées au calcul pour maximiser l'efficacité.

3. Contributions clés

Stratégie de parallélisme novatrice : Introduction du TSP, qui effondre deux axes de parallélisme orthogonaux en un seul, permettant la réduction simultanée de la mémoire des paramètres et des activations.
Analyse théorique : Un modèle mathématique rigoureux comparant le TSP au DP, TP, SP et TP+SP concernant :
- Mémoire : Le TSP est le seul schéma à réduire à la fois la mémoire proportionnelle aux poids et la mémoire des activations d'un facteur $1/D$ sans exiger $D^2$ dispositifs.
- Communication : Bien que le TSP introduise une communication de déplacement de poids (que le TP+SP évite), les auteurs prouvent que pour de longues séquences ou des tailles de lot modérées, la communication proportionnelle aux activations domine, rendant le TSP compétitif.
- FLOPs : Le TSP maintient la même efficacité de calcul par dispositif que le TP et le SP.
Implémentation et superposition : Implémentation détaillée des blocs d'attention et MLP du TSP avec des stratégies explicites de superposition communication/calcul (mise en file d'attente des transferts de poids derrière les calculs GEMM).
Optimisation topologique : Démonstration que le TSP permet de faire tenir l'ensemble du réplique du modèle dans un seul nœud à haute bande passante (intra-nœud), évitant les pénalités de performance liées à l'extension sur des liens inter-nœuds plus lents souvent requises par les grilles 2D TP+SP.

4. Résultats

Les auteurs ont évalué le TSP sur un cluster allant jusqu'à 1 024 GPU AMD MI300X (128 nœuds) avec un modèle transformeur dense de 7 milliards de paramètres.

Efficacité mémoire :
- Le TSP a constamment atteint la mémoire de pointe la plus faible sur toutes les longueurs de séquence testées (16K à 128K jetons).
- Pour des contextes courts, le TSP correspondait au TP (dominé par la mémoire des paramètres).
- Pour des contextes longs, le TSP a nettement surpassé le TP et le TP+SP (dominé par la mémoire des activations), empêchant les erreurs de mémoire insuffisante (OOM) là où d'autres stratégies échouaient.
Débit :
- Malgré un volume de communication théorique plus élevé (dû au déplacement de poids), le TSP a maintenu un débit compétitif avec le TP+SP.
- La communication supplémentaire a été masquée avec succès derrière le calcul via la mise en file d'attente.
- Les avantages de débit ont augmenté avec les degrés de parallélisme croissants ( $D$ ) et les longueurs de séquence.
Évolutivité :
- Le TSP a permis des micro-lots plus grands grâce à une empreinte mémoire réduite, améliorant davantage le débit.
- Il a permis de faire tenir les répliques de modèle entièrement dans un seul nœud (8 GPU), libérant des rangs pour le parallélisme de données (DP) dans des clusters plus grands.

5. Importance

Mise à l'échelle consciente du matériel : Le TSP fournit une solution pratique pour les clusters HPC modernes où la bande passante intra-nœud est élevée mais où la bande passante inter-nœud constitue un goulot d'étranglement. Il maximise l'utilité des interconnexions intra-nœud coûteuses.
Faisabilité des longs contextes : Alors que les modèles se dirigent vers des fenêtres de contexte plus longues (100k+ jetons), la mémoire des activations devient le goulot d'étranglement principal. Le TSP est unique pour gérer cela sans la surcharge massive de dispositifs du parallélisme 2D.
Flexibilité : Le TSP peut être composé de manière orthogonale avec d'autres schémas de parallélisme (parallélisme de pipeline, parallélisme d'experts, parallélisme de données), offrant un « axe » flexible pour que les concepteurs de systèmes optimisent pour des topologies matérielles spécifiques.
Efficacité des coûts : En réduisant le nombre de dispositifs nécessaires pour le parallélisme de modèle, le TSP permet d'utiliser plus de dispositifs pour le parallélisme de données, réduisant potentiellement la taille totale du cluster requise pour une tâche d'entraînement donnée ou augmentant la taille de lot effective.

En conclusion, le TSP représente une évolution significative des stratégies d'entraînement distribué, échangeant une augmentation gérable du volume de communication contre des gains substantiels en efficacité mémoire et en utilisation du matériel, en faisant un outil essentiel pour l'entraînement des LLM de nouvelle génération à longs contextes.