Ai2 ersetzt Prioritäten-Scheduler durch GPU-Zeitbudgets, um wirkungsvolle Forschung zu finanzieren

Nachfrage weit über dem Angebot

Das AI-Infrastructure-Team von Ai2 verwaltet Tausende NVIDIA-GPUs vom Typ H100, B200 und B300, die in Clustern mit 88 bis 1024 GPUs organisiert sind. Die Cluster sind für verteiltes Training im großen Maßstab ausgelegt und werden von rund 150 internen Forschenden genutzt. Deren Arbeit umfasst den gesamten Modell-Workflow aus LLM- und VLM-Training, Simulationen für Reinforcement Learning (RL) in der Robotik sowie das Post-Training für wissenschaftliche agentische Anwendungsfälle.

Wie in vielen Laboren übersteigt die Nachfrage das Angebot. Anhand der eingereichten Workloads beziffert Ai2 die offenen Anfragen auf das Zwei- bis Dreifache der jeweils verfügbaren GPUs. Um jede verfügbare GPU-Stunde konkurrieren also zwei bis drei Forschungs-Workloads.

Das Team beschreibt seine Aufgabe als Pyramide aus vier Kennzahlen. Die Basis bildet die Verfügbarkeit, darauf folgen die Auslastung der Belegung (Occupancy) und die Wirkung (Impact), als Spitze steht die Utilization. Der Beitrag konzentriert sich auf die Wirkung, definiert als die Häufigkeit, mit der die wertvollsten Workloads Ressourcen zugeteilt bekommen.

Was bei den Prioritäten schieflief

Der alte Scheduler arbeitete prioritätsbasiert, und Workloads konnten sich gegen Preemption entscheiden. Jedes Team hatte ein Limit für gleichzeitig genutzte GPUs für Workloads, die vor Preemption geschützt waren, während präemptible Workloads dieses Limit auf freien GPUs überschreiten durften.

Ai2 beschreibt mehrere vorhersehbare Probleme:

• GPU-Squatting: Nutzer hielten Leerlauf-Workloads vor, auf die sie sich bei Bedarf aufschalten konnten, weil sie Debugging-Workloads nicht mit ausreichend niedriger Latenz starten konnten, um Probleme in Echtzeit zu bearbeiten.
• Prioritätsinflation: Schließlich nutzten 100 Prozent der geplanten Workloads die Priorität HIGH, sodass niedrigere Prioritätsstufen keine GPU-Zeit mehr bekamen.
• Support-Aufwand: Weil die Präemptibilität optional war, verbrachten Bereitschaftsingenieure den Großteil ihrer Ticket-Bearbeitungszeit damit, das geordnete Herunterfahren nicht präemptibler Workloads auf Hosts mit bekannten Wartungsproblemen auszuhandeln.

Ai2 räumt ein, die Ursachen nur langsam erkannt zu haben. Frühe Maßnahmen konzentrierten sich auf eine strengere Kontrolle der Prioritäten und auf GPU-Monopole für wichtige Projekte. Heute sieht das Team die Lage als Tragödie der Allmende: Einzelne konkurrieren um eine knappe gemeinsame Ressource und erzeugen ein global suboptimales Ergebnis. Es verweist zudem auf das Paper „Dominant Resource Fairness“ von Ghodsi et al. aus dem Jahr 2011, das von einem Suchmaschinenunternehmen berichtet, dessen Nutzer „ihren Code mit Endlosschleifen spickten, um die Auslastung künstlich aufzublähen“.

Monopole waren Ai2 zufolge zu grob. Weil Teams zu unterschiedlichen Zeitpunkten bereit für Experimente sind, standen dedizierte GPUs still, während andere Teams auf Kapazität warteten.

Budgets und hierarchisches Fair-Share

Statt Teams GPUs zuzuweisen, vergibt Ai2 nun einen Anteil an GPU-Zeit. Die Begründung: Der künftige Bedarf hängt von den Ergebnissen neuartiger Experimente ab und lässt sich nicht exakt prognostizieren, während die Priorität zwischen Forschungsvorhaben eine strategische Frage ist, über die sich im Voraus debattieren lässt. Die Leitung soll wie Investoren denken und vor der Existenz konkreter Workloads anhand der voraussichtlichen Wirkung entscheiden, wie jedes Vorhaben gefördert wird.

Manager können GPU-Zeit anteilig den Projekten und Forschenden zuweisen, für die sie verantwortlich sind. In Ai2s Beispieldiagramm hat Projekt A1 einen Anspruch auf 35 Prozent der Gesamtkapazität, unabhängig davon, wie viele andere Projekte in der Warteschlange stehen. Jede Anfrage nach GPU-Zeit muss durch ein Budget gedeckt sein, sonst ist sie nicht vor Preemption geschützt. Ein Squatting-Workload gibt in den Worten von Ai2 Teambudget für nichts aus. Ziel ist es, das Austricksen des Schedulers teurer zu machen, als ehrlich für ein größeres Budget zu argumentieren.

Zuteilungsentscheidungen innerhalb eines Projekts trifft ein leitender Forscher, innerhalb eines Programms ein Principal Investigator und über Programme hinweg ein Lead Program Manager oder der CEO.

Die tatsächliche Belegung steuert ein hierarchischer Fair-Share-Scheduler. Ai2 betont, der Algorithmus sei nicht neu, und führt ihn auf den Hadoop Fair Scheduler von 2009 sowie heute auf den Fair Tree von SLURM und den Fair Scheduler von YARN zurück. Neu seien die Eingaben: Der Baum spiegelt die Struktur der Forschungsprogramme wider, und die Gewichte sind von Managern festgelegte Budgets statt statischer Quoten. Der Scheduler erfasst die Belegung über ein gleitendes Rückblickfenster, standardmäßig 7 Tage, und reiht Workloads aus unterausgelasteten Zuteilungen vor Workloads aus überausgelasteten ein.

Was das für die Forschenden bedeutet

Ai2 erwartet, dass jede Gruppe ihre zugeteilte GPU-Zeit über einen Zeitraum von einer Woche erhält, solange sie weiterhin Workloads mit ausreichendem Bedarf einreicht. Der Beitrag zitiert Chris Clark: „Der neue Scheduler fühlt sich an, als hätten wir 30 Prozent mehr Rechenleistung.“ Er sagt, stoßartige Workloads könnten die Zuteilungsgrenze nun später überschreiten und trotzdem schnell und ohne Preemption eingeplant werden. So könne das Team Rechenleistung zurückgewinnen, die der alte Scheduler faktisch verschenkt habe.

Der Wechsel verlagert zudem eine umstrittene Entscheidung aus dem Tagesbetrieb heraus. Laut dem Beitrag feilt Ai2 weiter am Budget-Review-Prozess: mit häufigen Gelegenheiten für Forschende, ihr Anliegen vorzutragen, und Entscheidungen durch Manager, die den meisten Kontext haben.