Résumé Technique : DOPPLER – Apprentissage à Double Politique pour l'Assignation de Dispositifs dans les Graphes de Flux de Données Asynchrones

1. Énoncé du Problème

Le document traite du problème d'assignation de dispositifs dans les systèmes multi-GPU exécutant des charges de travail complexes d'apprentissage automatique (ex: Transformers, Mixture-of-Experts) représentées sous forme de graphes de flux de données statiques. L'objectif est de mapper les opérations (sommets) sur des dispositifs afin de minimiser le temps d'exécution total au sein d'un système Work-Conserving (WC - de type conservation de travail).

Contrairement aux frameworks traditionnels à synchronisation globale (ex: PyTorch ou TensorFlow standard) où tous les dispositifs attendent l'opération la plus lente avant de procéder, les systèmes WC planifient les opérations de manière asynchrone dès que leurs dépendances de données sont satisfaites. Bien que les systèmes WC offrent un potentiel de gain de vitesse significatif en chevauchant communication et calcul, ils introduisent des défis uniques :

• Ordre d'Exécution Stochastique : Sans synchronisation globale, l'ordre d'exécution est dynamique et sensible à l'hétérogénéité du matériel et à la contention des ressources.
• Optimisation Complexe : En raison de cette stochasticité, le temps d'exécution ne peut pas être exprimé comme une fonction de forme fermée de l'assignation.
• Objectifs Concurrents : Une assignation efficace doit simultanément équilibrer la charge des GPU et minimiser la communication inter-GPU, une tâche rendue difficile car l'équilibrage de charge est intrinsèquement temporel dans les contextes asynchrones.

Les approches existantes basées sur l'apprentissage sont limitées par leur dépendance aux hypothèses de synchronisation globale, l'utilisation d'une politique de placement unique sans modélisation de la dynamique du système, et l'absence d'optimisation lors du déploiement réel.

2. Méthodologie : Le Framework DOPPLER

Les auteurs proposent DOPPLER, un framework d'apprentissage par renforcement (RL) qui emploie un schéma de décision séquentiel à double politique et un pipeline d'entraînement en trois étapes.

2.1 Formulation à Double Politique

Au lieu d'apprendre une politique monolithique unique, DOPPLER décompose le processus d'assignation en deux politiques distinctes, modélisées comme un processus de décision de Markov (MDP) épisodique :

1. SEL (Politique de Sélection) : Sélectionne la prochaine opération (sommet) à assigner parmi l'ensemble des nœuds candidats. Cette politique parcourt le graphe partiellement assigné pour approximer le flux non déterministe du "temps", apprenant ainsi efficacement un ordre d'exécution.
2. PLC (Politique de Placement) : Détermine le dispositif GPU spécifique pour l'opération sélectionnée, visant à équilibrer la charge et à minimiser la communication.

Architecture :

• Les deux politiques utilisent des Réseaux de Neurones sur Graphes (GNN) avec passage de messages pour encoder la structure du graphe de flux de données, les dépendances et les chemins critiques.
• SEL agrège les informations des chemins critiques (chemins de bas niveau et de haut niveau) et les caractéristiques des nœuds pour sélectionner le nœud suivant.
• PLC combine les représentations des nœuds, les caractéristiques des dispositifs (ex: charge actuelle, temps de fin) et l'ensemble des nœuds déjà placés pour sélectionner le dispositif optimal.
• Optimisation de l'Efficacité : Pour éviter les coûts d'entraînement prohibitifs associés à la réexécution du passage de messages à chaque étape, DOPPLER effectue le passage de messages une seule fois par épisode et met à jour les caractéristiques des dispositifs dynamiquement sans re-calculer les plongements (embeddings) du graphe à chaque étape.

2.2 Pipeline d'Entraînement en Trois Étapes

DOPPLER emploie un paradigme d'« apprentissage par l'action » (learning-by-doing) à travers trois étapes pour assurer la convergence et l'adaptabilité :

1. Étape I (Apprentissage par Imitation / Hors-ligne) : Les politiques sont pré-entraînées par apprentissage supervisé pour imiter un enseignant heuristique (spécifiquement, l'heuristique du Chemin Critique). Cela accélère la convergence en fournissant une politique initiale forte avant que l'exploration ne commence.
2. Étape II (RL basé sur la Simulation / Hors-ligne) : Les politiques sont affinées via l'apprentissage par renforcement au sein d'un simulateur logiciel (implémentant la logique d'exécution WC). La récompense est dérivée du temps d'exécution simulé. Cette étape permet une itération rapide sans la latence du matériel réel.
3. Étape III (RL sur Système Réel / En ligne) : Les politiques entraînées sont déployées dans un système WC réel. Elles sont continuellement affinées via du RL en ligne en utilisant les temps d'exécution réels comme récompenses. Cette étape adapte la politique aux caractéristiques matérielles spécifiques (ex: jitter de NVLink, contention) que les simulateurs ne peuvent capturer parfaitement, sans surcharge supplémentaire pour l'utilisateur.

3. Contributions Clés

1. Formulation du Problème : Le papier étudie l'assignation de dispositifs spécifiquement dans le contexte des systèmes multi-GPU asynchrones et work-conserving, traitant la nature stochastique de l'ordre d'exécution.
2. Approche à Double Politique : Il introduit une séparation des préoccupations où une politique apprend l'ordre de parcours des nœuds tandis que l'autre gère le placement des dispositifs, capturant ainsi à la fois la dynamique d'exécution et les contraintes matérielles.
3. Framework en Trois Étapes : DOPPLER propose une stratégie d'entraînement hybride combinant l'apprentissage par imitation, le RL basé sur la simulation et le RL en ligne sur système réel. Cette approche accélère la convergence et permet une optimisation continue lors du déploiement.
4. Implémentation GNN Scalable : Les auteurs proposent une approximation efficace du passage de messages qui réduit considérablement le temps d'entraînement sans compromettre la convergence.

4. Résultats Expérimentaux

Les expériences ont été menées sur quatre GPU NVIDIA Tesla P100 en utilisant diverses charges de travail : multiplications de matrices en chaîne (CHAINMM), réseaux de neurones feedforward (FFNN) et blocs/couches de transformateurs Llama.

• Performance vs Baselines : DOPPLER (spécifiquement le DOPPLER-SYS entièrement entraîné) a systématiquement surpassé les bases de comparaison, incluant l'heuristique du Chemin Critique, PLACETO (RL à politique unique), GDP (RL par plongement de graphe) et un optimiseur énumératif personnalisé.
  • DOP_PLER a atteint des temps d'exécution jusqu'à 52,7 % inférieurs par rapport à la meilleure base (GDP).
  • Il a réduit le temps d'exécution jusqu'à 78,2 % par rapport à l'heuristique du Chemin Critique.
  • Il a surpassé l'optimiseur énumératif personnalisé de près de 13,8 %.
• Études d'Ablation :
  • Double Politique : La combinaison des politiques SEL et PLC a produit les meilleurs résultats, particulièrement pour les modèles complexes, démontant la nécessité d'apprendre à la fois l'ordre et le placement.
  • Étapes d'Entraînement : L'entraînement sur le système réel seul a conduit à une convergence lente et instable. La combinaison de l'apprentissage par imitation (Étape I) et de la simulation (Étape II) a été cruciale pour une convergence stable et rapide avant le raffinement en ligne (Étape III).
• Transfer Learning : DOPPLER a démontré une forte capacité de généralisation. Des politiques entraînées sur des graphes simples (FFNN) pouvaient être affinées avec peu d'épisodes (2k–4k) pour performer sur des architectures inédites (Llama) et différentes configurations matérielles (ex: transfert de 4 P100 vers 8 V100).
• Scalabilité : Les temps d'entraînement et d'inférence de DOPPLER évoluent linéairement avec la taille du graphe de flux de données, surpassant les autres bases de comparaison RL en termes d'efficacité.

5. Signification et Revendications

Le papier affirme que DOPPLER représente une avancée significative dans l'optimisation des charges de travail ML modernes en dépassant les hypothèses statiques de synchronisation globale. Sa signification réside dans :

• Efficacité Pratique : En permettant une exécution work-conserving, DOPPLER peut économiser des heures de GPU substantielles dans les déploiements à grande échelle (ex: économisant des millions d'heures de GPU annuellement pour des charges de travail de l'échelle de ChatGPT).
• Adaptabilité : Contrairement aux heuristiques statiques ou aux modèles purement basés sur la simulation, la capacité d'apprentissage en ligne de DOPPLER lui permet de s'adapter aux réalités bruitées des environnements matériels de production.
• Innovation Méthodologique : La décomposition en double politique offre une approche plus nuancée du problème d'explosion combinatoire de l'assignation de dispositifs, en séparant le "quoi exécuter ensuite" du "où l'exécuter", ce qui s'aligne mieux avec la nature temporelle de l'exécution asynchrone.

Les auteurs soulignent que leur approche ne nécessite pas de modifications des frameworks de ML sous-jacents, mais optimise plutôt la couche de planification (scheduling), ce qui en fait une solution pratique pour les infrastructures multi-GPU existantes.