Ai2 zmienia scheduler GPU: budżety czasu zamiast priorytetów

Artykuł źródłowy opublikowany na blogu Hugging Face opisuje, jak zespół infrastruktury AI w Allen Institute for AI (Ai2) zmienił sposób przydzielania czasu na klastrach GPU swoim badaczom. Zamiast kolejnego dopracowywania systemu priorytetów, inżynierowie Ai2 zdecydowali się przebudować model od podstaw — zastępując klasyczny scheduler priorytetowy systemem budżetów czasu GPU, hierarchicznej sprawiedliwej alokacji (fair-share) oraz kontraktu time-slicingowego.

Tekst, napisany przez Kyle’a Wiggersa z zespołu Ai2Comms, tłumaczy, dlaczego dotychczasowe podejście prowadziło do patologii znanych z ekonomii zasobów współdzielonych, i jak zmiana filozofii — z „planowania zadań” na „budżetowanie czasu” — pozwoliła ograniczyć te problemy.

Kluczowe wnioski

• Ai2 zarządza tysiącami GPU Nvidia H100, B200 i B300, zorganizowanych w klastry od 88 do 1024 jednostek, obsługujące ok. 150 wewnętrznych badaczy.
• Popyt na GPU przewyższa dostępną podaż 2–3 razy — na każdą dostępną godzinę GPU przypadają 2–3 konkurujące o nią zadania badawcze.
• Stary system priorytetowy doprowadził do „squattingu” GPU, inflacji priorytetów (ostatecznie 100% zadań miało priorytet HIGH) oraz sytuacji, w której inżynierowie dyżurni większość czasu spędzali na negocjowaniu wyłączania niepreemptowalnych zadań.
• Nowy model opiera się na przydzielaniu zespołom nie konkretnych GPU, lecz portfela czasu GPU, co pozwala zachować pełne obłożenie klastra niezależnie od sezonowości potrzeb badawczych.
• Zespół odwołuje się do klasycznego problemu „tragedii wspólnego pastwiska” oraz do pracy naukowej z 2011 roku o Dominant Resource Fairness, by pokazać, że problem alokacji zasobów nie jest unikalny dla AI.

Piramida czterech metryk i problem z priorytetami

Zespół infrastruktury AI w Ai2 opisuje swoje zadanie jako piramidę czterech nadrzędnych metryk, które budują się jedna na drugiej. Fundamentem jest dostępność — jak często sprzęt jest sprawny i gotowy do pracy. Wyżej znajduje się obłożenie, czyli odsetek dostępnego czasu przypisanego do konkretnego zadania. Trzeci poziom to wpływ (impact) — jak często zasoby trafiają do najbardziej wartościowych projektów badawczych. Szczytem piramidy jest wykorzystanie (utilization) — faktyczny odsetek mocy GPU zużyty w trakcie trwania zadania.

Artykuł koncentruje się na drugim od góry poziomie, czyli na poprawie trafności decyzji schedulera co do tego, które projekty otrzymują GPU. Historycznie Ai2 korzystało ze schedulera opartego na priorytetach, który pozwalał zespołom „wypisywać się” z preemptowalności (czyli możliwości przerwania zadania na rzecz innego, ważniejszego). Każdy zespół miał limit liczby GPU chronionych przed preemption, a zadania preemptowalne mogły przekraczać ten limit na bezczynnych maszynach.

System ten generował przewidywalne patologie. Badacze zaczęli „parkować” bezużyteczne zadania (no-op) na GPU tylko po to, by mieć do nich dostęp w razie potrzeby debugowania w czasie rzeczywistym — bo uruchomienie nowego zadania z odpowiednio niskim opóźnieniem było niemożliwe. Z czasem doszło też do inflacji priorytetów: ostatecznie praktycznie wszystkie zgłaszane zadania miały priorytet HIGH, co oznaczało, że niższe poziomy priorytetów były całkowicie pozbawione dostępu do GPU. Dodatkowo, ponieważ preemptowalność była opcjonalna, dyżurujący inżynierowie większość czasu reakcji na zgłoszenia spędzali na negocjowaniu zorganizowanego wyłączenia niepreemptowalnych zadań działających na maszynach z problemami serwisowymi.

Tragedia wspólnego pastwiska i inspiracja ekonomią zasobów

Zespół Ai2 przyznaje, że początkowo nie rozpoznał źródła tych problemów. Pierwsze próby naprawy polegały na zaostrzeniu kontroli nad ustalaniem priorytetów, a ostatecznie na obejściu schedulera poprzez ręczne przypisywanie „monopoli” na GPU najważniejszym projektom. Jak zauważają autorzy, firma nieświadomie zbudowała idealne „laboratorium” do obserwacji tragedii wspólnego pastwiska — sytuacji, w której jednostki konkurujące o ograniczony, współdzielony zasób, maksymalizując własne korzyści, prowadzą do nieoptymalnego wyniku globalnego i nadużywania zasobu.

Autorzy przypominają, że problem ten nie jest nowy w badaniach nad alokacją zasobów, łączących algorytmikę, ekonomię i zarządzanie systemami. Centralnym problemem jest fakt, że użytkownicy często lepiej znają wartość swoich własnych zadań niż organizacja, ale mają motywację, by tę wartość skrywać lub zatrzymywać zasoby nawet wtedy, gdy szkodzi to ogólnej wydajności. Jako przykład przywołują pracę Ghodsiego i współautorów z 2011 roku wprowadzającą koncepcję Dominant Resource Fairness — opisuje ona anegdotę o firmie z branży wyszukiwania, która przydzielała dedykowane maszyny tylko tym zadaniom, które gwarantowały wysokie wykorzystanie zasobów. Użytkownicy szybko odkryli, że mogą „nasycić” swój kod nieskończonymi pętlami, by sztucznie zawyżyć wskaźniki wykorzystania. Sprzęt się zmienia, ale fundamentalne problemy czyniące alokację zasobów trudną — pozostają takie same.

Od monopoli GPU do budżetów czasu

Klasycznym rozwiązaniem tragedii wspólnego pastwiska jest prywatyzacja współdzielonego zasobu — właściciele mają wtedy motywację do maksymalizacji jego wartości. Ai2 już wcześniej próbowało wersji tego podejścia, przypisując zespołom monopole na zestawy GPU, ale okazało się to zbyt grubą kreską: powodowało bezczynność sprzętu wynikającą z sezonowości badań. Zespoły są gotowe do uruchamiania eksperymentów i treningów w różnych momentach, więc przypisanie monopolu gwarantowało okresy, w których żadne zadania nie były gotowe do wykonania, a inny zespół czekał na dostęp do mocy obliczeniowej.

W praktyce zespół ręcznie rozwiązywał problem przypominający klasyczny problem plecakowy (knapsack problem) — starając się wpasować dynamicznie zmieniające się potrzeby badawcze w statyczny harmonogram. Autorzy chcieli zachować motywację „własnościową”, jednocześnie utrzymując pełne obłożenie GPU.

Rozwiązaniem okazała się iteracja modelu własności: zamiast przydzielać zespołom konkretne GPU, Ai2 zaczęło alokować im porcję czasu GPU. Przewidywanie przyszłego popytu wymagałoby znajomości wyników jeszcze nieprzeprowadzonych eksperymentów naukowych, więc nie da się go precyzyjnie prognozować. Priorytet pomiędzy projektami badawczymi jest jednak kwestią strategii, którą można łatwiej debatować i ustalać z wyprzedzeniem. Zamiast próbować rozwiązać zagadkę harmonogramowania, system umożliwił kierownictwu myślenie w kategoriach inwestora — decydowanie, jeszcze przed powstaniem konkretnych zadań, jak sfinansować każdy projekt badawczy w czasie GPU, na podstawie ich oceny wartości strategicznej.

Tego typu wyzwania związane z rozdzielaniem ograniczonej infrastruktury obliczeniowej są dziś powszechne w branży AI, w miarę jak rośnie zapotrzebowanie na infrastrukturę GPU niezbędną do trenowania coraz większych modeli.

Podsumowując, zespół Ai2 przekształcił debatę o tym, ile czasu GPU zasługuje każdy projekt badawczy, z codziennej, operacyjnej improwizacji w przejrzysty, administracyjny proces budżetowania — podejście, które może stać się inspiracją dla innych laboratoriów AI zmagających się z chronicznym niedoborem mocy obliczeniowej.