Google atvėrė TorchTPU kodą darbui su savo DI lustais

„Google“ viešai atvėrė „TorchTPU“ kodą, leidžiantį „PyTorch“ modelius vykdyti jos TPU procesoriuose ir išlaikyti įprastas šios programavimo sistemos sąsajas. Spalio 20 dienos „Google Cloud“ pranešime projektas pristatomas kaip bendras darbas su „Meta“ komanda, mažinantis modelių perkėlimo į kitokią aparatinę įrangą kliūtis.

Projekto GitHub saugykloje nurodyta „Apache 2.0“ licencija ir pateikiamas atsisiunčiamas kodas, dokumentacija bei modelių pavyzdžiai savarankiškam bandymui. Tai nėra naujas kalbos modelis ar mokama pokalbių paslauga: „TorchTPU“ yra programinis sluoksnis, padedantis esamus „PyTorch“ skaičiavimus vykdyti konkrečioje procesorių šeimoje.

Pasak saugyklos aprašo, bandymams reikia TPU virtualiosios mašinos, o kaip paprastesnis pradinis variantas įvardijama „Google Cloud“ v6e instancija. Šiuo metu palaikoma „Python 3.11“ ir naujesnė versija; vien įprasto nešiojamojo kompiuterio be prieigos prie TPU šiems skaičiavimams nepakanka, nors kodas viešai prieinamas.

„Hugging Face“ savo tos pačios dienos pristatyme aiškina, kad ankstesnis kelias dažnai vedė per „PyTorch/XLA“ arba reikalavo mokytis „JAX“. „TorchTPU“ įtraukia TPU kaip dar vieną įrenginio tipą, todėl modelių kūrėjai gali likti prie pažįstamos „PyTorch“ darbo eigos ir bibliotekų.

„Google Cloud“ nurodo, kad įrankis palaiko tiesioginio vykdymo režimą derinimui ir kompiliavimą gamybiniam našumui, taip pat paskirstytus skaičiavimus. Integracijos apima „Transformers“, „Diffusers“, „vLLM“ ir „SGLang“, tačiau skirtingų modelių palaikymą bei našumą kūrėjai turi vertinti pagal konkrečias bibliotekas ir jų dokumentaciją.

„Hugging Face“ skelbia, kad iš TPU vykdomų jos testų praeina 96,5 procento, o iš 34 praneštų problemų 19 jau ištaisytos. Šie skaičiai apibūdina bibliotekos bandymų būklę, o ne visų pasaulyje egzistuojančių „PyTorch“ programų suderinamumą, todėl leidėjai patys pažymi dar likusias nesėkmes ir darbų sąrašą.

Tame pačiame „Hugging Face“ pristatyme pateikiamas našumo pavyzdys: 64 užklausų grupėje specialus branduolys „Qwen3-30B-A3B“ iškodavimą paspartino 21,5 karto prieš numatytąją realizaciją. Rezultatas gautas su TPU v6e-8, todėl jo negalima tiesiogiai laikyti palyginimu su „Nvidia“ įranga ar pažadu kiekvienam modeliui ir apkrovai.

Saugykloje pateikiami „Llama“, „Qwen“, „ResNet“, DLRM ir kitų architektūrų pavyzdžiai, taip pat derinimo ir našumo matavimo priemonės. Dokumentacijoje nurodomas duomenų bei tenzorių lygiagretinimo palaikymas, todėl projektas skirtas ne tik vieno modelio demonstracijai, bet ir skirtingiems paskirstyto darbo scenarijams su TPU.

„Hugging Face“ pabrėžia, kad ankstesnis „Optimum TPU“ tarpinis sprendimas nebereikalingas modeliams su naująjį palaikymą turinčiomis „Transformers“ ir „Diffusers“ bibliotekomis vykdyti. Kūrėjai tuo pat metu nurodo tęsiamus testus, todėl naujas būdas nepanaikina poreikio patikrinti konkretaus modelio veikimą ir naudojamų bibliotekų versijų suderinamumą.

Saugyklos apraše šiuo metu nurodytas viešas „0.2.0“ leidimas, o „Google“ pranešime aprašomos tolesnės mokymo, modelių aptarnavimo ir diagnostikos integracijos. „Hugging Face“ planuoja tęsti reguliarius TPU testus ir problemas perduoti projekto kūrėjams, todėl artimiausią pažangą rodys nauji pataisymai bei suderinamumo rezultatai.