Az Ai2 új GPU-ütemezővel osztja el a szűkös számítási időt

Több GPU-ra lenne szükség, mint amennyi rendelkezésre áll

Az Ai2 mesterségesintelligencia-infrastruktúráért felelős csapata több ezer NVIDIA H100, B200 és B300 GPU-t kezel. A fürtök mérete 88 és 1024 GPU között változik, és elsősorban nagy léptékű, elosztott modelltréningre készültek.

A kapacitást nagyjából 150 belső kutató használja. Munkájuk az LLM-ek és VLM-ek teljes tréningfolyamatát, a robotikai megerősítéses tanulás szimulációját, valamint tudományos ügynöki felhasználási esetek utótréningjét is lefedi. Az Ai2 szerint a beadott feladatok alapján bármely pillanatban a rendelkezésre állónál 2-3-szor több GPU-ra érkezik igény.

A szervezet a kapacitáskezelést négy egymásra épülő mutatóval írja le. Az alap a rendelkezésre állás, vagyis hogy a hardver milyen gyakran egészséges és munkára kész. Erre épül a foglaltság, amely azt mutatja meg, a rendelkezésre álló idő mekkora részét rendelik egy feladathoz. Következik a hatás, vagyis milyen gyakran jutnak erőforráshoz a legértékesebb munkák, a csúcs pedig a kihasználtság, a GPU-kapacitás feladat élettartama alatti tényleges használata.

A prioritásalapú rendszer több problémát is okozott

A korábbi ütemező prioritások alapján döntött, és lehetővé tette, hogy a feladatok kivonják magukat a megszakíthatóság alól. Az egyes csapatok korlátozott számú GPU-t használhattak ilyen védett munkákhoz, a megszakítható feladatok pedig az üres kapacitást is igénybe vehették.

Az Ai2 több visszatérő problémát figyelt meg. Egyes felhasználók üres, úgynevezett GPU-parkoló feladatokat indítottak, amelyekhez később csatlakozhattak. Erre azért volt szükségük, mert a hibakeresési feladatokat nem tudták elég gyorsan elindítani, amikor valós időben felmerült egy probléma.

Emellett kialakult a prioritások felduzzadása is. Idővel minden ütemezett feladat HIGH prioritást használt, így az alacsonyabb szintekhez tartozó munkák teljesen kiszorultak. Az ügyeletes mérnökök idejük jelentős részét azzal töltötték, hogy egyeztessék a nem megszakítható feladatok leállítását olyan gépeken, amelyeken karbantartási problémát észleltek.

Az Ai2 ezt a közös erőforrások tragédiájaként értelmezi. Az egyéni érdekek követése ebben a helyzetben a teljes rendszer számára kedvezőtlen eredményhez és az erőforrás helytelen használatához vezethet.

A GPU-k helyett a használati idő kap költségkeretet

Az Ai2 korábban egyes projekteknek lényegében GPU-monopóliumot adott, ez azonban üresen hagyott kapacitáshoz vezethetett. A kutatási igények időben változnak, ezért előfordulhatott, hogy egy csapat nem készült futtatásra, miközben másik csapat várt a kapacitásra.

Az új megközelítésben a csapatok nem konkrét GPU-kat kapnak, hanem a GPU-idő egy részére jogosító költségkeretet. A szervezet vezetői előre dönthetnek arról, mely kutatási irányokat milyen arányban finanszírozzák GPU-idővel, a döntést pedig a projektek várható hatásáról alkotott megítélésükre alapozhatják.

A hierarchikus rendszerben a vezetők arányosan oszthatják el a GPU-időt a hozzájuk tartozó projektek és kutatók között. Az Ai2 példája szerint az A1 projekt a teljes kapacitás 35 százalékára tarthat igényt, függetlenül attól, hogy máshol hány projekt várakozik.

Minden GPU-időre vonatkozó igényt költségkeretnek kell fedeznie, különben a feladat nem kap védelmet a megszakítással szemben. Így az üresen parkoltatott munka is a felhasználó keretét terheli. Az Ai2 célja, hogy az ütemező kijátszása drágább legyen, mint a nagyobb keretért folytatott nyílt egyeztetés.

Hét napos időablak alapján osztják el a kapacitást

A költségkeretek mellé az Ai2 hierarchikus méltányos részesedési ütemezőt épített. Az algoritmus a szervezeti programstruktúrát követi, a súlyokat pedig a vezetők által meghatározott költségkeretek adják.

Az ütemező gördülő időablakban követi a foglaltságot, amelynek alapértelmezett hossza hét nap. Az alulkihasznált keretekhez tartozó feladatokat az ütemező előrébb sorolja a túlkihasznált keretek feladatainál. Az Ai2 szerint így egyhetes időtávon minden csoport megkaphatja a neki kiosztott GPU-időt, ha folyamatosan elegendő feladatot küld be.

Az új rendszer gyakorlati előnyeiről Chris Clark, az Ai2 munkatársa azt mondta: „Az új ütemezővel úgy érezzük, mintha 30 százalékkal több számítási kapacitásunk lenne.” Szerinte korábban elveszett az a kapacitás, amelyet egy csapat időnként nem használt ki, most viszont a későbbi időszakokban az elosztás fölé kerülhetnek, és a feladataik gyorsan, megszakítás nélkül indulhatnak.

A változás a kutatók számára azt jelenti, hogy a GPU-idő elosztása kevésbé eseti üzemeltetési kérdésként, inkább átlátható költségkeret-alapú döntésként jelenik meg. A keretekről a megfelelő szervezeti szinten döntenek, a vezető kutatótól a programvezetőn át akár a vezérigazgatóig.