Construire un ordonnanceur de cluster pour prioriser la recherche à fort impact tout en maintenant une occupation complète
Au sein de l'équipe AI Infrastructure d'Ai2, nous sommes responsables de la fourniture de la capacité de calcul GPU de l'institut, en ciblant spécifiquement les charges de travail d'entraînement distribuées à grande échelle. Nous concevons cette tâche comme une pyramide de quatre métriques qui se construisent les unes sur les autres.
La fondation est la disponibilité: la fréquence à laquelle le matériel est sain et prêt à travailler. Au-dessus se trouve l'occupation: la fraction du temps disponible assignée à une charge de travail spécifique. Ensuite vient l'impact: la fréquence à laquelle les charges de travail les plus précieuses sont choisies pour recevoir des ressources. Le sommet de la pyramide est l'utilisation: la fraction de la capacité GPU utilisée sur la durée de vie d'une charge de travail.
Cet article porte sur l'amélioration de l'impact de nos décisions d'ordonnancement. Nous avons récemment remplacé un ordonnanceur basé sur les priorités par un système incluant des budgets de temps GPU, une allocation hiérarchique en parts équitables et un contrat de découpage temporel. En conséquence, nous avons transformé le débat sur le temps GPU que mérite chaque projet de recherche, d'une tâche opérationnelle au cas par cas en un processus transparent de budgétisation administrative.
Sur-engagement
Chez Ai2, nous gérons des milliers de GPU NVIDIA H100, B200 et B300, répartis dans des clusters dont la taille varie de 88 à 1024 GPU. Ces clusters sont conçus pour l'entraînement distribué à grande échelle de modèles d'IA, et ils servent un groupe d'environ 150 chercheurs internes dont le travail couvre un ensemble diversifié de domaines de l'IA, incluant le flux complet des modèles pour l'entraînement de LLM et de VLM, la simulation d'apprentissage par renforcement (RL) en robotique, et le post-entraînement pour des cas d'usage agentiques scientifiques.
Comme beaucoup de laboratoires, nous avons une demande de temps GPU qui dépasse largement l'offre. D'après les charges de travail soumises, à tout moment nous avons des demandes en attente pour 2 à 3 fois plus de GPU qu'il n'y en a de disponibles. Une façon de le voir est que chaque heure GPU disponible sur notre cluster fait l'objet d'une compétition entre 2 à 3 charges de travail de recherche différentes.
Historiquement, nous utilisions un ordonnanceur basé sur les priorités, et nous permettions aux charges de travail de se soustraire à la préemption. Chaque équipe avait une limite de GPU simultanés pouvant être utilisés par des charges de travail protégées contre la préemption. Les charges de travail préemptibles pouvaient dépasser cette limite sur des GPU inactifs. Cette stratégie produisait des pathologies prévisibles. Par exemple, nous avons observé des cas d'« occupation illicite » (squatting) de GPU où des utilisateurs laissaient tourner des charges de travail sans effet auxquelles ils pouvaient se connecter au besoin. Cela se produisait parce que les chercheurs constataient qu'ils ne pouvaient pas lancer des charges de travail de débogage avec une latence assez faible pour traiter les problèmes en temps réel. Nous avons aussi observé une inflation des priorités, où finalement 100% des charges de travail ordonnancées utilisaient la priorité HIGH. Cela signifiait que les niveaux de priorité inférieurs étaient complètement privés de temps GPU. Comme la préemption était optionnelle, nous avons également constaté que nos ingénieurs d'astreinte passaient la majorité de leur temps de réponse aux tickets à négocier l'arrêt organisé de charges de travail non préemptables s'exécutant sur des hôtes présentant des problèmes de maintenance connus.
Tragédie des biens communs
Lorsque ces problèmes sont apparus, nous avons été lents à en identifier les causes profondes. Nos premières tentatives pour garantir que le travail le plus important reçoive du temps GPU se concentraient sur un contrôle plus strict de la manière dont les priorités étaient définies et, finalement, sur le contournement de l'ordonnanceur basé sur les priorités en assignant explicitement des monopoles GPU à des projets importants. Bien que nous ne l'ayons pas reconnu au début, nous avions construit un laboratoire parfait pour observer la « tragédie des biens communs ». Des individus se disputaient une ressource partagée rare et, en cherchant à maximiser leurs résultats individuels, obtiennent un résultat global non optimal et abusent de la ressource sous-jacente.
Nous étions loin d'être les premiers à observer ce type d'interaction. L'allocation des ressources est un domaine de recherche fascinant qui mêle développement d'algorithmes, économie et gestion de systèmes. Un problème central est que les utilisateurs connaissent souvent mieux que l'organisation la valeur de leurs propres travaux, mais ils peuvent avoir des incitations à masquer cette valeur ou à conserver des ressources même lorsque cela nuit à la performance globale. Par exemple, dans leur article de 2011 introduisant Dominant Resource Fairness, Ghodsi et al. racontent une anecdote dans laquelle une société de recherche sur le web fournissait des machines dédiées à des travaux uniquement si leurs utilisateurs pouvaient garantir une haute utilisation. Ils ont rapidement découvert que « les utilisateurs parsemaient leur code de boucles infinies pour gonfler artificiellement les niveaux d'utilisation ». Le matériel change, mais les problèmes fondamentaux qui rendent complexe l'allocation des ressources persistent.
Des budgets plutôt que des ordonnancements
La solution classique à une tragédie des biens communs est de privatiser la ressource partagée—les propriétaires sont incités à maximiser la valeur de leur bien. Lorsque nous attribuions aux équipes des monopoles sur des ensembles de GPU, nous faisions déjà une version de cela, mais c'était trop grossier. Cela entraînait des GPU inactifs en raison de la saisonnalité de la recherche. Les équipes sont prêtes à lancer des expériences et des entraînements à des moments différents, donc attribuer un monopole garantissait qu'il y aurait des périodes où aucun travail n'était prêt à s'exécuter, pendant qu'une autre équipe attendait de la capacité.
Nous résolvions manuellement un problème de sac à dos (knapsack problem)en essayant de faire tenir des besoins de recherche changeant dynamiquement dans un calendrier statique. Nous voulions l'incitation à l'appropriation, mais nous voulions aussi maintenir une occupation complète des GPU.
Nous avons décidé d'itérer sur le modèle d'appropriation. Au lieu d'attribuer des GPU aux équipes, nous avons choisi d'allouer une portion de temps GPU. Prédire la demande dans le futur exigerait de connaître le résultat d'expériences scientifiques inédites, ce qui ne peut donc pas être prévu avec précision. La priorité entre les efforts de recherche, en revanche, est une question de stratégie, et elle peut être plus facilement débattue et décidée à l'avance. Au lieu d'essayer de résoudre le casse-tête de l'ordonnancement, nous avons permis au leadership de penser comme des investisseurs. Avant que les charges de travail n'existent, décider comment financer chaque effort de recherche en temps GPU en fonction de leur jugement de son impact probable. L'ordonnanceur pouvait ensuite utiliser cette information lors de la priorisation des charges de travail entrantes.
Avec cela en tête, nous avons conçu un système hiérarchique où les managers pouvaient allouer proportionnellement du temps GPU aux projets et aux chercheurs dont ils étaient responsables. Comme l'illustre le schéma ci-dessous, cela traduit directement la stratégie du programme en une part garantie de temps GPU. Le projet A1 sait qu'il détient une revendication de 35% sur la capacité totale, quel que soit le nombre d'autres projets en file d'attente ailleurs.
Les valeurs entre parenthèses représentent la capacité totale du cluster attribuée à un projet feuille.
Dans ce système, chaque demande de temps GPU doit être financée par un budget, sinon elle n'est pas protégée contre la préemption. Dans l'ancien système, la priorité HIGH n'avait aucun coût et la non-préemptibilité permettait à une équipe de remplir indéfiniment sa limite de GPU concurrents, donc tout le monde les utilisait. Désormais, rien n'est gratuit, donc toute astuce pour obtenir du temps GPU puise dans l'allocation de l'utilisateur qui en bénéficie. Une charge de travail qui monopolise des ressources dépense le budget de l'équipe pour rien. Notre stratégie est de rendre la manipulation de l'ordonnanceur plus coûteuse que la participation honnête au débat pour un budget plus important. Nous itérons constamment sur ce processus de revue du budget, mais les exigences clés sont qu'il existe des occasions fréquentes pour les chercheurs de plaider pour le temps dont ils ont besoin, et que les décisions soient prises par les managers ayant le plus de contexte sur les arbitrages en question. Cela signifie que les décisions d'allocation au sein d'un projet de recherche sont prises par un chercheur principal, au sein d'un programme de recherche par un investigateur principal, et entre les programmes par un responsable de programme principal, ou par le PDG.
Part équitable (fair-share)
Associé à cet outil de budgétisation du temps GPU, nous avons construit un ordonnanceur hiérarchique à part équitable pour gérer l'occupation effective des allocations throughout l'arbre du programme. L'algorithme ici n'est pas nouveau—la part équitable hiérarchique sur une fenêtre temporelle fait partie d'une lignée qui remonte au Hadoop Fair Scheduler en 2009, et la même approche est activement utilisée aujourd'hui dans SLURM's Fair Tree et YARN's Fair Scheduler. Ce qui est nouveau pour nous, ce sont les entrées : l'arbre reflète la structure du programme de recherche, et les poids sont des budgets fixés par les managers plutôt que des quotas statiques.
L'ordonnanceur suit l'occupation sur une fenêtre glissante de rétrospection (7 jours par défaut) et trie les charges de travail des allocations sous-utilisées au-dessus de celles des allocations sur-utilisées. De cette façon, sur une plage d'une semaine, nous pouvons nous attendre à ce que chaque groupe reçoive son temps GPU alloué tant qu'il soumet activement des charges de travail avec une demande suffisante.
« Le nouvel ordonnanceur donne l'impression que nous avons 30% de calcul en plus. Avec l'ancien ordonnanceur, si nous avions des moments où nous n'utilisions pas toute notre limite de slots, ce calcul était pratiquement perdu. Maintenant, avec le nouvel ordonnanceur, si cela se produit, nous pouvons ensuite dépasser notre limite d'allocation et voir nos tâches toujours ordonnancées rapidement et sans préemption, ce qui nous permet essentiellement de récupérer ce calcul. Nos charges de travail sont souvent en rafales, cela nous a donc rendu une quantité significative de calcul. » — Chris Clark
L'ordonnanceur distingue deux types d'occupation. L'occupation allouée est le temps pendant lequel une charge de travail est imputée à un budget. Cela puise dans les allocations du propriétaire de la charge de travail, ce qui affecte le calcul du budget à part équitable, et ces charges de travail sont protégées contre la préemption pendant leur fenêtre d'exécution minimale. L'occupation non allouée n'est imputée à aucun budget, n'est pas protégée dès le départ, et peut être préemptée par toute demande allouée. Cela nous permet de garder les GPU pleinement occupés même lorsque les allocations ne correspondent pas correctement à la demande et empêche les équipes de jamais refuser des cycles GPU gratuits.
Le contrat d'ordonnancement
Une caractéristique supplémentaire de l'entraînement distribué qui rend difficile une allocation équitable des ressources est que les charges de travail peuvent fonctionner très longtemps. Les tâches d'entraînement s'exécutent régulièrement pendant des heures, des jours, et parfois même des semaines. Une fois ordonnancée, une charge de travail peut rester sur ses GPU assignés pendant une semaine ou plus, ne laissant aucune opportunité aux autres de recevoir leur temps budgété. C'est la propriété du système qui a rendu possible la monopolisation des GPU. C'est aussi ce qui obligeait les ingénieurs d'astreinte à négocier avec les propriétaires de tâches de longue durée pour résoudre les problèmes de maintenance en cours.
Pour résoudre ces problèmes, nous avons introduit un « contrat d'ordonnancement ». En échange de l'accès au cluster, une charge de travail doit déclarer son temps d'exécution minimal, c'est-à-dire la durée d'occupation la plus courte nécessaire pour réaliser des progrès significatifs. Pendant cette période, la charge de travail est protégée de la préemption. Cela donne au chercheur une garantie de progression, tout en donnant à l'ordonnanceur le droit de rééquilibrer une fois cette progression acquise, en remettant automatiquement en file d'attente les charges de travail reprenables. Alternativement, un utilisateur peut définir un temps d'exécution minimal de zéro, ce qui indique que le temps GPU doit être désalloué. Ces charges de travail sont toujours soumises à la préemption, mais elles sont également gratuites en ce sens qu'elles ne sont pas facturées à aucun budget.
Le cycle de vie d'une charge de travail suit ce schéma :
- La charge de travail est soumise avec un temps d'exécution minimal et indique si elle est reprenable ou non.
- La charge de travail est ordonnancée selon l'algorithme de partage équitable, pondéré par un rapport entre l'occupation réelle et le temps alloué dans la fenêtre rétrospective.
- La charge de travail s'exécute pendant son temps d'exécution minimal, qui est facturé à ses allocations.
- La charge de travail peut continuer à s'exécuter tant que les allocations associées continuent de la prioriser par rapport aux autres. Ce temps est également facturé à ses allocations.
- Elle peut être préemptée et remise en file d'attente, ce qui ramène à l'étape 2.
- La charge de travail se termine, libérant sa revendication sur toutes les ressources.
Ensemble, ces accords ajoutent le découpage temporel à notre ordonnanceur. Les charges de travail en cours d'exécution peuvent être retirées et remises en file d'attente automatiquement, permettant au partage équitable de converger et décourageant l'occupation abusive. Elles permettent également aux hôtes défaillants de vider leurs charges de travail lorsqu'elles atteignent leur temps d'exécution minimal, de sorte que les activités de réparation peuvent être entièrement automatisées. Ce dernier point s'est révélé plus important que nous ne l'avions réalisé en planifiant ce travail. Il a réduit de 74 % les réparations nécessitant une intervention humaine, ce qui représentait une économie considérable de corvée d'astreinte.
Simulations
Nous savons que les changements de politique d'ordonnancement peuvent avoir des conséquences involontaires. La nature à somme nulle du problème signifie que donner du temps à un chercheur revient à en retirer à un autre. Les utilisateurs qui perdent cet échange ont tendance à chercher de nouveaux contournements. Avant de déployer le système basé sur les budgets, nous voulions un moyen rapide de prédire où ces temps d'attente plus longs pourraient apparaître et de tester des paramètres de configuration comme la longueur de la fenêtre rétrospective ou la valeur maximale autorisée pour le temps d'exécution minimal (nous avons choisi 8 heures).
Nous avons construit un petit environnement de simulation qui prend en entrée un ensemble de charges de travail et leur calendrier de soumission et permet à l'ordonnanceur de prendre des décisions de préemption et d'affectation de GPU. Avec la connaissance du nombre de GPU demandé et de la durée totale d'exécution de chaque charge de travail, le simulateur pouvait avancer jusqu'aux moments ordonnançables et fournir une analyse des temps d'attente en file d'attente, des événements de préemption et de la répartition du temps GPU entre les projets pour de nombreux jours simulés en quelques secondes. Nous avons exécuté le simulateur à la fois sur des données de soumission historiques et sur des scénarios construits que nous voulions mieux comprendre.
Une hypothèse que nous voulions tester concernait les « charges de travail de débogage ». Ces tâches nécessitent un petit nombre de GPU et un temps d'exécution minimal de 15 minutes ou moins, ce qui suffit à l'utilisateur pour voir si une tâche se lance avec succès ou plante rapidement en raison d'un bug ou d'une mauvaise configuration. Nous voulions savoir si ces tâches connaîtraient un temps d'attente en file plus court que les charges de travail d'entraînement plus importantes, qui ont souvent besoin de nombreux GPU et d'heures d'exécution pour réaliser des progrès significatifs. Intuitivement, ces tâches plus petites devraient remonter en tête de file, car une petite tâche peut tenir à plus d'endroits qu'une grande. Mais la latence précise de la file était importante. Une attente courte d'une ou deux minutes débloquerait une nouvelle pratique de développement, mais une attente de dix minutes devient irréalisable.
Nos simulations nécessitaient des données de cas de test construites à la main, car notre enregistrement historique ne contenait pas un volume suffisant de ces charges de travail de type débogage. Nos résultats ont soutenu l'hypothèse, montrant que les temps d'attente p90 des charges de travail de débogage passent d'environ 6 heures à seulement 5 minutes.
Visualisation du simulateur à plus petite échelle du scénario de référence (à gauche) et du nouvel ordonnanceur « allocations » (à droite). Chaque ligne est un GPU ; chaque barre est une tâche, colorée par charge de travail parente avec une teinte par équipe ; le hachurage marque le temps où une tâche est interruptible, et un bord rouge marque une préemption. Dans le scénario de référence, les longues tâches urgentes ne sont jamais interrompues, et moins de préemptions se produisent sur les travaux de priorité inférieure. Le nouvel ordonnanceur présente un mélange plus important de couleurs sur chaque GPU, illustrant la rotation d'occupation entre les équipes.
Résultats
Avec les résultats de simulation en main, nous avons commencé un déploiement cluster par cluster fin juillet. Les résultats qui nous importent sont de savoir si les charges de travail que nous avions choisi de financer ont reçu leur temps, si le nouveau système a maintenu une occupation complète, et si les chercheurs pouvaient raisonner sur l'ordonnanceur pour prendre des décisions éclairées.
Depuis le déploiement, nous avons observé que les utilisateurs et les équipes reçoivent systématiquement leur temps GPU alloué. Nous comptons le temps dû à une équipe comme son plafonnement d'allocation heure par heure à sa demande réelle. Sur la période de test de 30 jours, les équipes ont reçu 98 % des heures GPU qui leur étaient dues, et 13 des 15 allocations d'équipe ont reçu 95 % ou plus, le pire cas recevant 90 %. L'occupation du cluster s'est maintenue stable à 98 % avant et après le changement, avec une demande dépassant la capacité d'un facteur 2-3x dans les deux périodes. 18 % du temps GPU fourni était non alloué, ce qui nous a permis de maintenir une occupation élevée pendant les périodes où les cas d'usage financés n'étaient pas prêts à s'exécuter.
Les résultats de notre simulateur se sont révélés directionnellement exacts, les résultats réels dépassant nos prévisions. Le temps d'attente p90 dans la file des charges de travail de debug est passé de 2 heures à 30 secondes sous le nouvel ordonnanceur, contre une prédiction simulée de 6 heures à 5 minutes issue de scénarios de test élaborés manuellement. Il convient de noter que la taille plus réduite de l'échantillon des charges de travail de debug dans la baseline impliquait une variance plus élevée dans ces mesures. La latence de file en général s'est améliorée comme effet secondaire du time-slicing : sur notre plus grand cluster H100, le temps d'attente médian dans la file est passé de 5 minutes à 24 secondes, et le temps d'attente p90 a diminué d'environ un tiers (de 2.8 heures à 1.8 heures).
Par rapport aux trois problèmes que nous cherchions à résoudre :
Le squatting : Les courtes charges de travail de debug démarrent en moins d'une minute, ce qui réduit l'intérêt du squatting. Le coût de ce comportement est débité du budget du squatter, ce qui l'empêche de recevoir du temps lorsqu'il en a vraiment besoin.
L'inflation des priorités : Nous permettons toujours aux charges de travail de déclarer une priorité, mais celle-ci n'affecte que le tri au sein d'une équipe. Les managers sont incités à surveiller les priorités du groupe afin d'optimiser l'utilisation de leurs budgets.
La corvée d'astreinte : Les hôtes malsains sont drainés automatiquement lorsque les charges de travail atteignent leur temps d'exécution minimum. Les réparations nécessitant une intervention humaine ont diminué de 74 %.
Défis
La courbe d'apprentissage a été plus raide que nous ne l'avions supposé. Nous avons déployé le changement de manière incrémentale, donc durant les premiers jours, les chercheurs rencontraient des comportements variés selon le cluster qu'ils ciblaient. De plus, nos interfaces conservaient une ancienne terminologie (comme la priorité des charges de travail) dont le sens avait changé. La documentation seule n'a pas résolu la confusion. Ce qui a fonctionné fut l'organisation de sessions explicatives en direct, offrant un forum où les chercheurs pouvaient poser des questions et où l'équipe d'ingénierie pouvait fournir des descriptions plus approfondies de la manière et des raisons pour lesquelles l'ordonnanceur prenait ses décisions de priorisation, à l'aide d'exemples réels.
Ce fut un moment clé car il a marqué un tournant après une période initiale de frustration et de théories populaires vers le mode actuel, où les groupes de recherche communiquent plus souvent et plus largement sur les besoins en GPU de leurs expériences. Les chercheurs participent désormais à la discussion sur les budgets avec une meilleure connaissance des compromis faits pour répondre à toute nouvelle demande.
En plus des sessions en personne, nous avons introduit de nouvelles visualisations après le lancement pour donner aux utilisateurs une meilleure idée de la proximité entre leur temps GPU alloué et leurs allocations attendues, et pour exposer directement la métrique utilisée pour trier la file des charges de travail. Cela a fourni un endroit simple où regarder lorsqu'une charge de travail était préemptée pour comprendre pourquoi. Ces visualisations ont également aidé les propriétaires de budgets, qui pouvaient voir comment le temps GPU était utilisé dans les différents projets sous leur gestion.
Exemple de visualisation de l'utilisation de l'allocation dans le temps.
Tous les cas d'usage ne se sont pas améliorés. En plus de l'entraînement distribué, nos chercheurs lancent des sessions interactives où ils effectuent des analyses de données et testent leur code d'entraînement pendant qu'ils l'écrivent. Dans l'ancien système, un chercheur pouvait maintenir une telle session jusqu'à une semaine. Avec le time-slicing, ils étaient soumis à la limite de 8 heures du runtime protégé, après quoi une session devenait préemptible si elle dépassait son allocation. Nous n'avions pas mesuré à quel point les chercheurs dépendaient de l'état volatile de ces sessions. Être préempté signifiait attendre pour obtenir une nouvelle session et aussi reconstruire leur état à la main. Après avoir interrogé les chercheurs pour comprendre l'ampleur de ce problème, nous avons créé deux nouveaux projets de feuille de route. Nous investissons dans un cluster CPU uniquement, à côté de notre stockage on-prem, pour des sessions de développement axées sur les tâches de préparation de données. Cela préservera la capacité de notre cluster d'entraînement pour les charges de travail qui en ont véritablement besoin. De plus, nous prévoyons de construire des sessions restaurables pour ces charges de travail CPU uniquement. Cela nous permettra de continuer à préempter les charges de travail à la fin de leur temps d'exécution minimum pour la maintenance ou le time-slicing, tout en pouvant restaurer la session ailleurs sans que le chercheur ait à la reconstruire. Nous pouvons conserver les bénéfices opérationnels et d'ordonnancement de ce nouveau système tout en améliorant l'expérience utilisateur.
Nous restons à l'affût des problèmes émergents. Un problème potentiel que nous étudions est la fragmentation de la capacité, qui peut entraîner une augmentation des temps d'attente dans la file pour les plus grandes charges de travail. Notre intuition est que la protection par temps d'exécution minimum s'applique aux types de tâches qui comptaient auparavant sur des mécanismes préemptibles pour dépasser la limite de GPU concurrents de leur équipe. Auparavant, ces tâches pouvaient être interrompues à tout moment, gaspillant potentiellement ce temps, mais facilitant aussi l'ordonnancement des grandes tâches. Désormais, l'ordonnanceur peut avoir moins d'occasions d'interrompre de nombreuses tâches à la fois pour placer une grande charge de travail en attente. Nous utilisons actuellement nos outils de simulation pour reproduire ce problème tout en mesurant la vérité terrain en production.
L'avenir
Au-delà du travail d'ordonnancement décrit ici, nous visons le sommet de la pyramide : l'utilisation. Nous devons garantir que le bootstrapping, le checkpointing et les applications d'entraînement elles-mêmes soient tous exécutés aussi efficacement que possible, maximisant la valeur du temps ordonnancé que chaque charge de travail reçoit.
Si vous souhaitez relever des défis comme ceux-ci en collaboration étroite avec des chercheurs, nous vous encourageons à explorer les postes d'ingénierie ouverts chez Ai2.