A PyTorch natív eszközként kezeli az IBM Spyre gyorsítót

A PyTorch rétegeihez kapcsolják a Spyre-t

A PyTorch október 8-án közzétett bemutatója szerint a torch-spyre a keretrendszer meglévő eszköz-, memóriafoglaló-, adatfolyam- és fordítóabsztrakcióit kapcsolja össze a Spyre futtatókörnyezetével és firmware-ével. A PyTorch így a spyre azonosítón keresztül kezelheti a gyorsítót, miközben a kártyával közvetlenül a Spyre futtatókörnyezete és firmware-e kommunikál.

Az integráció a PrivateUse1 regisztrációra épül. A PyTorchban átnevezik ezt a háttér-szolgáltatást spyre-ra, majd regisztrálják a hozzá tartozó eszközmodult. A megfelelő memóriafoglaló, másolási műveletek és diszpécser-regisztrációk mellett így működhet a tensor.to("spyre") művelet, a rendszer pedig a Spyre-hoz tartozó kernelekhez irányítja a tenzorokat.

A fejlesztők a backend saját vezérlőit a torch.spyre alatt érhetik el, míg a torch.accelerator hordozható felületet ad az elérhetőség, az eszközválasztás, az adatfolyamok, a szinkronizáció és a memóriajelentések kezeléséhez.

A hardver felépítése meghatározza a működést

A Spyre kártyán 32 mag kapcsolódik nagy sávszélességű gyűrűn keresztül. Minden maghoz 2 MB helyi tárterület és feldolgozóelemekből álló tömb tartozik, a tenzorok és programok tárolására pedig legfeljebb 128 GB LPDDR5 memória áll rendelkezésre.

A két memóriaszint feladata eltérő. Az LPDDR5-öt a futtatókörnyezeti verem kezeli, a fordító pedig előre előállítja azokat a betöltési és tárolási műveleteket, amelyek a csempéket a memória és az egyes magok helyi tárterülete között mozgatják. Az adatmozgatás 128 bájtos egységekben történik, ezért az elrendezés és az igazítás hatással van a betöltés hatékonyságára.

A Spyre számítási modellje lefordított programokra épül. Ezek több funkcionális egység, például feldolgozóelemek, speciális függvénycsatornák és betöltési, tárolási egységek programjait tartalmazzák. A hardveres sorok és a helyi memóriák kötik össze ezeket az egységeket, a haladás pedig akkor történik, amikor a fogyasztó egység számára elérhetővé válnak a szükséges operandusok.

Adatfolyamokkal és eseményekkel kezelik az átfedést

A Spyre külön számítási és adatmozgatási csatornákat használ. Így egy adatátvitel párhuzamosan futhat egy lefordított program számításával. A külön beadott számítási műveletek ugyanakkor egy közös futtatókörnyezeti számítási sorban várakoznak egymás mögött, ezért a futtatókörnyezeti átfedés az adatmozgatás és a számítás párhuzamosításából származik, nem két egyidejű számítási indításából.

A PyTorch adatfolyamai biztosítják az egyes műveletek sorrendjét. A rendszer a mozgatási és számítási munkát külön adatfolyamokhoz rendelheti, majd eseménnyel kötheti össze azokat, amikor az egyik művelet a másik eredményét használja. Egyetlen adatfolyam használata megőrzi a helyességet, de sorosítja a különálló csatornákat.

A torch.compile(backend="inductor") továbbra is a PyTorch Inductor fordítási útvonalán tartja az FX-gráfokat. A lefordított modellek és a regisztrált, azonnali műveletek ugyanazt a fordítási, gyorsítótárazási és indítási útvonalat használhatják. A lefordított program indításakor a rendszer a PyTorch által átadott tenzorokhoz köti a futáskor ismert címeket, miközben a programhoz tartozó köztes és ideiglenes tárhelyek címei előre rögzíthetők.

Mit jelent ez a fejlesztőknek?

A megközelítés célja, hogy a fejlesztők ismerős PyTorch-kódon keresztül használhassák az inferenciára optimalizált Spyre-t. A PyTorch kezeli a tenzorok élettartamát is: amikor egy Spyre-tenzorra már nem marad hivatkozás, a keretrendszer visszahívja a backend memóriafoglalóját, amely felszabadíthatja vagy újrahasznosíthatja az eszközön lefoglalt területet.

A PyTorch közleménye szerint az eredmény egy közös útvonal az azonnali és a lefordított végrehajtáshoz, kisebb indítási többlettel. A Spyre-t az IBM olyan vállalati csapatok számára tervezte, amelyek AI-feladatokat futtatnak alkalmazásaik és adataik mellett IBM Z, LinuxONE és Power rendszereken. A gyorsító csökkentett pontosságú számításai a nyelvgenerálásban és a beágyazási modellek mátrixigényes műveleteihez illeszkednek.