Résumé Technique : DYNATRAIN

Énoncé du Problème

L'entraînement moderne des grands modèles de langage (LLM) repose sur des compositions complexes de stratégies de parallélisme (partage des données, des tenseurs, du pipeline, des experts et de l'optimiseur) pour adapter les modèles à des clusters GPU massifs. Cependant, les frameworks d'entraînement actuels fonctionnent sous un modèle d'exécution statique, supposant que la disposition du parallélisme choisie au démarrage reste optimale pendant toute la durée de vie de la tâche. Cette hypothèse devient de plus en plus intenable en raison de :

1. Fluctuations des Ressources : Les clusters de production réallouent dynamiquement les ressources en raison de la préemption, de la révocation des instances spot, ou du fait que les charges de travail de service co-localisées libèrent des GPU pendant les heures creuses.
2. Flux de Travail Dynamiques : Les flux de travail émergents comme l'apprentissage par renforcement à partir de retours humains (RLHF) nécessitent de déplacer des ressources entre les phases d'inférence (déroulement) et d'entraînement, rendant nécessaires des changements de parallélisme à l'exécution.

Les solutions existantes présentent des limitations significatives :

• Sauvegarde et Redémarrage : L'approche standard consiste à arrêter la tâche, à sérialiser des téraoctets d'état sur disque, puis à redémarrer. Cela entraîne une latence massive (de dizaines de secondes à plusieurs minutes) due aux goulots d'étranglement des E/S de stockage et au coût élevé de la réinitialisation des communicateurs distribués.
• Systèmes Élastiques Existants : Les approches actuelles de « commutation à chaud » en mémoire sont souvent limitées à des modèles prédéfinis, à des dimensions spécifiques de parallélisme ou à des topologies statiques, échouant à prendre en charge une reconfiguration arbitraire multidimensionnelle.

Le défi central pour permettre une commutation à chaud générale est triple :

1. Hétérogénéité de l'État : Les optimisations modernes (par exemple ZeRO) fragmentent les états selon plusieurs dimensions, créant des shards disjoints et hétérogènes qui ne s'alignent pas sur les limites des tenseurs.
2. Routage Complexe : La transition entre deux dispositions multidimensionnelles est un problème de communication pair-à-pair de type M-vers-N, nécessitant une planification sans blocage mutuel sans vue globale du changement de topologie.
3. Contraintes Mémoire : Les anciens et nouveaux états doivent coexister temporairement avec les tampons de communication, risquant des échecs de mémoire insuffisante (OOM) sur des GPU à capacité limitée.

Méthodologie

DYNATRAIN est un système d'entraînement distribué conçu pour permettre une reconfiguration en ligne en moins d'une seconde à travers des stratégies de parallélisme arbitraires. Il fonctionne comme une couche middleware enfichable entre le backend d'apprentissage profond (par exemple PyTorch) et les frameworks d'entraînement (par exemple Megatron-LM).

1. Espace de Paramètres Virtuel (VPS)

L'abstraction centrale de DYNATRAIN est l'Espace de Paramètres Virtuel (VPS), un espace de coordonnées logiques unifié où tous les paramètres, gradients et états d'optimiseur existent sous leur forme complète, non fragmentée.

• Représentation Logique : Chaque tenseur est représenté par une boîte englobante d'index multidimensionnelle logique. Les dimensions parallèles (DP, TP, PP, EP, ZeRO) sont réduites à une taille de 1 dans cet espace.
• Mappage Déterministe : Toute configuration de parallélisme agit comme une fonction de mappage mathématique ($F_P$) qui projette le VPS global vers un sous-VPS local à chaque rang.
• Intersection Géométrique : La reconfiguration entre stratégies se réduit au calcul de l'intersection géométrique entre les régions de sous-VPS source et destination. Cela transforme la migration complexe d'état en un problème déterministe d'identification des régions logiques à envoyer, recevoir ou conserver, évitant l'explosion des règles de conversion par paires.

2. Composants du Système

Construit sur le VPS, DYNATRAIN utilise trois composants coopératifs :

• Planificateur de Routage d'État : Opère entièrement dans les coordonnées du VPS. Il compare les régions logiques source ($R^{(i)}_{src}$) et destination ($R^{(i)}_{dst}$) pour chaque rang afin de dériver les ensembles exacts d'envoi, de réception et de conservation. Il gère le désalignement des états d'optimiseur (qui sont souvent partitionnés par des décalages de tampon plats plutôt que par des limites de tenseurs sémantiques) en les mappant à nouveau vers les coordonnées du VPS, les traitant comme des intersections géométriques standard.
• Moteur de Transition d'État : Affine le plan de haut niveau du planificateur en un pipeline d'exécution optimisé.
  • Optimisation des Primitives : Il promeut les transferts point-à-point vers des primitives collectives (diffusion, dispersion, collecte) lorsque la topologie le permet, et sélectionne les paires source-destination optimales en fonction de la bande passante intra-nœud (par exemple NVLink).
  • Mise en Mémoire Consciente : Pour prévenir les OOM, il effectue la collecte de mémoire des tampons obsolètes et partitionne les transferts en étapes en fonction de la mémoire disponible. Il utilise des tampons d'envoi/réception contigus pour fusionner les transferts fragmentés, maximisant ainsi l'utilisation de la bande passante.
  • Planification Sans Blocage Mutuel : Il emploie un algorithme de planification asynchrone basé sur le XOR ($Peer(i, s) = i \oplus s$) pour garantir des paires de communication symétriques, assurant une séquence sans blocage mutuel sans coordination globale.
• Gestionnaire de Périphériques Élastique (EDM) : Découple les stratégies logiques des groupes de communication physiques.
  • Il gère le cycle de vie des groupes de communication, les mettant en cache pour une réutilisation instantanée.
  • Il prend en charge le chevauchement asynchrone, permettant à la boucle d'entraînement de continuer sur l'ancienne topologie tandis que le nouveau « Monde » et les sous-groupes sont initialisés en arrière-plan, masquant ainsi la latence des changements de topologie.

Contributions Clés

1. Espace de Paramètres Virtuel (VPS) : Une nouvelle abstraction qui unifie les états distribués hétérogènes sous un seul système de coordonnées logiques, réduisant les transitions de parallélisme à des problèmes d'intersection géométrique.
2. Mécanismes Cohésifs de Commutation à Chaud : Une suite d'outils incluant le Planificateur de Routage d'État, un Moteur de Transition d'État conscient de la mémoire avec une planification prouvée sans blocage mutuel, et un Gestionnaire de Périphériques Élastique qui chevauche la reconfiguration avec l'entraînement.
3. Système de Bout en Bout : DYNATRAIN intègre ces mécanismes pour fournir une reconfiguration en ligne en moins d'une seconde pour les modèles denses et MoE, prenant en charge des changements arbitraires dans les dimensions DP, TP, PP et EP.

Résultats Expérimentaux

Le système a été évalué sur un cluster GPU de 16 nœuds (NVIDIA A800) utilisant des modèles allant de 1,3 à 235 milliards de paramètres (incluant LLaMA-2, GPT-3 et Qwen-MoE).

• Vitesse de Reconfiguration :
  • Modèles Denses : Reconfigure un modèle dense de 70 milliards de paramètres en moins de 2 secondes.
  • Modèles MoE : Reconfigure un modèle MoE de 235 milliards de paramètres en 4,36 secondes.
• Gains de Performance :
  • Surpasse Megatron-LM Distributed Checkpointing (MCP) jusqu'à 143 fois pour les modèles denses et 926 fois pour les modèles MoE.
  • Surpasse Tenplex (système élastique basé sur la sauvegarde) jusqu'à 33 fois pour les changements de dimension unique et 15 fois pour la migration de tâches inter-clusters.
  • Surpasse HotSpa (référence en mémoire) jusqu'à 2,11 fois dans les scénarios de reconfiguration hybride.
• Correction : La validation par convergence de la perte sur un modèle de 13 milliards de paramètres a montré que la trajectoire de perte de DYNATRAIN suit parfaitement une référence statique sans divergence ni pics, même avec des augmentations et diminutions périodiques de ressources tous les 50 pas.
• Mise à l'Échelle Élastique : Dans les scénarios de mise à l'échelle dynamique du cluster (ajout/suppression de nœuds), le système masque plus de 93 % de la surcharge d'initialisation en la chevauchant avec l'entraînement en cours, réduisant l'interruption perçue à moins de 3 secondes.

Importance

L'article soutient que la commutation à chaud du parallélisme devrait être une primitive système fondamentale pour adapter l'entraînement aux ressources dynamiques et aux flux de travail multiphases. En découplant les états d'entraînement distribués de leurs dispositions physiques grâce au VPS, DYNATRAIN rend la reconfiguration en ligne pratique plutôt qu'exceptionnelle. Ce travail représente une étape vers des environnements d'exécution d'entraînement adaptatifs où le parallélisme n'est pas une décision d'infrastructure fixe, mais une dimension de contrôle flexible capable de répondre continuellement aux dynamiques de la charge de travail et du cluster.