A Google és a Meta bemutatta a PyTorch natív TPU-hátterét

A korábbi PyTorch TPU-megoldások korlátai

A Google TPU-i gépi tanulásra tervezett, specializált áramkörök. A forrás szerint nagy léptékű modellek betanítására és futtatására a Google-nél és azon kívül is sikeresen használják őket, teljesítmény-, energia- és költségelőnyöket kínálva.

A TPU-kat hagyományosan a TensorFlow és a JAX ökoszisztémájával társították. A Google 2020-ban mutatta be a PyTorch első fordított TPU-változatát, a Lazy Tensorokra épülő PyTorch/XLA-t. A PyTorch szerint ez az integráció nem tudott lépést tartani a keretrendszer közösség által vezérelt fejlődésével. A használatához több esetben módosítani kellett a forráskódot, jelentős fordítási késleltetéssel járt, nem támogatta a PyTorch natív, imperatív végrehajtási módját, és késve kezelte a PyTorch elosztott számítási primitívjeit.

A TorchTPU-t ezzel szemben kifejezetten PyTorch-háttérként építették fel. A projekt nem használ JAX-függőségeket a futtatókörnyezetben, kiemelt szerepet ad a közvetlen végrehajtásnak, és a torch.compile() rendszert a TorchDynamo alapján integrálja. A fejlesztők célja a PyTorch elosztott megoldásainak támogatása is, köztük a DDP, az FSDP és a DTensor.

Minimális módosítással indulhatnak a modellek

A TorchTPU egyik központi ígérete, hogy egy meglévő munkafolyamat TPU-ra állítása egyszerű legyen. A forrás példája szerint a model.to("cuda") hívás model.to("tpu") formára cserélhető. Ez önmagában nem garantálja a TPU maximális teljesítményét, mivel a modellek alakjai, felosztása és műveletei eredetileg GPU-architektúrához igazodhatnak, a modellnek azonban a projekt célja szerint azonnal futnia kell.

A fejlesztők a közvetlen és a fordított mód mellett egy új, DeferAndFuse nevű közvetlen végrehajtási módot is hozzáadtak. A TorchTPU jelenleg körülbelül 450 Aten-műveletet támogat, a strukturális vagy összetett delegálást igénylő műveletek kivételével. A projekt szerint a fő mélytanulási architektúrákhoz szükséges alapvető Aten-műveletek előre- és visszafelé irányuló végrehajtása is támogatott.

A csapat a Hugging Face 5000 legtöbbet letöltött modelljén ellenőrizte a közvetlen, módosítás nélküli futtatást. A felsorolt modellek között szerepel többek között a Qwen, a DeepSeek, a Gemma, a Llama, a Mistral, a GPT, a Whisper, a BERT, a T5, a CLIP, a ResNet és a ViT.

Elosztott futtatás és egyedi kernelek

A TorchTPU támogatja a DTensor-alapú felosztást, valamint a több gazdagépes következtetést a vLLM és az SGLang használatával. A projektet a forrás szerint a TorchTitan segítségével végzett finomhangolási és tanítási munkafolyamatokon is ellenőrizték, Ray-alapú vezérléssel. A validált integrációk között a vLLM, az SGLang, a TorchTitan, a Ray és a TorchRec szerepel.

Egyedi kernelek JAX/Pallas és a Helion DSL segítségével is készíthetők. A TorchTPU fejlesztői a futtatási és fordítási hibákhoz használható, érthetőbb üzeneteket, valamint a PyTorch-közösség által már használt eszközök, például a torch.profiler, a DTensor és a Ray integrációját is célként jelölték meg.

A projekt aktív fejlesztés alatt áll, és a bejelentés szerint egyes nagy Google Cloud-ügyfelek éles következtetési és tanítási munkafolyamataiban már használják. A TorchTPU nyílt forráskódú projektként érhető el a hozzá tartozó GitHub-tárhelyen.

A TPU-k működése és a teljesítmény kihasználása

A Google TPU-k TensorCore egységeket, mátrixszorzó egységeket, vektoros és skaláris egységeket, valamint egyes generációkban SparseCore-okat tartalmaznak. A TPU 8i 288 GB kapacitású, 8,6 TB/s sávszélességű nagy sávszélességű memóriát kínál a forrás szerint. A TPU 8t és a TPU 8i chipjei közötti kétirányú kapcsolat 19,2 TB/s, a TPU 8t szeletek közötti adatközponti hálózata pedig 400 GB/s sebességű.

A TPU-kat az XLA fordítói verem programozza, többek között a gépi tanulási műveletekhez használt StableHLO és az egyedi kernelek alacsonyabb szintű leírására szolgáló Mosaic révén. A teljesítmény kihasználásához fontos, hogy a modell alakjai kellően szélesek legyenek a mátrixszorzó egységek képességeinek kiaknázásához. A fordító emellett a ritka és kollektív műveleteket SparseCore egységekre helyezheti át.

A projekt egyik fejlesztési iránya a korlátozott dinamizmus támogatása. Ennek célja, hogy a fordító elkerülje az alakváltozások miatti újrafordításokat, miközben az alak maximális méretei fordításkor, a tényleges méretek pedig futás közben ismertek. Ez a PyTorch-munkafolyamatok TPU-s használatának egyszerűsítését szolgálja.