PyTorch konferencē Eiropā 20. oktobrī Google un Meta kopīgi publicēja TorchTPU. Tas ir jauns atklātā koda aizmugures dzinējs, kas ļauj PyTorch modeļiem strādāt tieši uz Google Tensor Processing Unit (TPU) čipiem. TPU ir Google projektēti procesori, kas optimizēti matricu operācijām un tiek izmantoti gan Google iekšējos servosos, gan Google Cloud klientiem. Iepriekšējais ceļš caur PyTorch/XLA bibliotēku prasīja JAX ekosistēmas starpslāni, kas ierobežoja izstrādātāju brīvību. TorchTPU šo atkarību novāc.
Projekts sākās Google iekšienē 2025. gada vidū. Pēc aptuveni pusotra gada tas ir nonācis līdz ražošanas gatavībai: Meta jau lieto TorchTPU reālām darba slodzēm. Google publicēja kodu GitHub platformā ar Apache 2.0 licenci.
Vecais ceļš caur JAX bija kompromiss
PyTorch ir populārākais ietvars dziļās mācīšanās modeļu apmācībai. Pēdējos gados tas ir kļuvis par standartu gan pētniecībā, gan ražošanā. Taču līdz šim tas labi darbojās tikai uz NVIDIA GPU ar CUDA ekosistēmu. Google TPU čipiem bija pieejama PyTorch/XLA bibliotēka, kas tulkoja PyTorch operācijas caur JAX izpildes vidi.
Problēma slēpās fundamentālā arhitektūras atšķirībā. JAX izmanto atliktas izpildes (lazy evaluation) modeli: operācijas vispirms tiek saliktas grafā, tad izpildītas kopā. PyTorch pēc noklusējuma izpilda katru operāciju uzreiz (eager execution). Šī atšķirība piespieda izstrādātājus pielāgot treniņu ciklus, mainīt datu plūsmu, reizēm pārrakstīt modeļa fragmentus. PyTorch emuāra ierakstā Google raksta, ka vecās bibliotēkas “piespieduši PyTorch strādāt JAX paradigmā”.
TorchTPU izvēlas pretēju pieeju. Dzinējs ir eager-first arhitektūra bez JAX atkarībām izpildes laikā. Izstrādātājs raksta parastu PyTorch kodu, palaiž to uz TPU. Starpslāņa nav.
Trīs izpildes režīmi
TorchTPU piedāvā trīs eager režīmus ar atšķirīgu kompromisu starp ērtībām un ātrumu.
Debug izpilda katru operāciju atsevišķi ar pilnu kļūdu izsekošanu. Režīms ir lēns, bet ļauj precīzi lokalizēt problēmu konkrētā koda vietā. Strict Eager izpilda operācijas tūlītēji bez optimizācijām. Fused Eager automātiski apvieno blakus esošas operācijas vienā izsaukumā, piemēram, sapludina vairākas matricu reizināšanas. Google publicētajos mērījumos Fused Eager ir par 50 līdz 100 procentiem ātrāks nekā Strict Eager tajā pašā aparatūrā.
Pilnai grafika kompilācijai TorchTPU izmanto Torch Dynamo, kas uztver FX grafu. Kompilācijai kalpo XLA aizmugure (PyTorch noklusējuma Torch Inductor netiek izmantots). Visiem režīmiem ir kopīga kompilācijas kešatmiņa, ko var iestatīt kā pastāvīgu daudzserveru konfigurācijās.
Viena praktiska nianse: pašreizējās paaudzes TPU čipi sasniedz maksimālu matricu reizināšanas efektivitāti pie dimensijām 128 vai 256, bet daudzi modeļi cietā kodā lieto uzmanības galvu (attention head) dimensiju 64. Google iesaka pielāgot šo parametru, lai pilnībā izmantotu aparatūras iespējas.
Meta jau apkalpo produkcijas trafiku
Meta ir līdzautors projektā. Uzņēmums jau apkalpo lielo valodas modeļu, reklāmu atlases un rekomendāciju darba slodzes uz TPU caur TorchTPU. Konkrētus veiktspējas skaitļus uzņēmums nav publiskojis, bet dzinēja lietošana produkcijā norāda, ka tas ir izgājis ārpus eksperimenta stadijas.
TorchTPU ir jauns mēģinājums, kas sākts Google 2025. gada vidū sadarbībā ar PyTorch komandu Meta. Tā mērķis: izveidot natīvu un principiālu PyTorch aizmuguri Google TPU čipiem. — PyTorch emuārs, 2026. gada 20. oktobris
Vārds “principiāls” (principled) šeit nozīmē, ka dzinējs seko PyTorch iekšējai loģikai. Tas nepiespiež lietotāju pielāgoties svešam izpildes modelim, kā tas bija ar PyTorch/XLA.
Ironwood čipi un Hugging Face atbalsts
Pirmie publicētie veiktspējas dati attiecas uz Google jaunāko Ironwood (TPU v7) arhitektūru. Google ziņo par 52 procentu modeļa FLOPS efektivitāti uz Qwen 3 8B modeļa ar gandrīz lineāru mērogošanu līdz 8192 Ironwood čipiem. Neatkarīgu verifikāciju šiem skaitļiem pagaidām nav, tāpēc tie jāuztver kā ražotāja pašnovērtējums.
Hugging Face Transformers v5 bibliotēkā integrācija jau darbojas. Pietiek norādīt device_map="tpu", lai modelis tiktu izvietots uz TPU. Iepriekš šāda pārslēgšanās prasīja atsevišķu konfigurāciju katrai aparatūrai ar vairākām koda rindām tikai ierīces norādīšanai.
CUDA gravitācija saglabājas
Dienu pirms Google paziņojuma Amazon publicēja TorchNeuron ar līdzīgu mērķi saviem Trainium čipiem. Abi projekti cenšas samazināt pārejas izmaksas no NVIDIA CUDA ekosistēmas, kas joprojām dominē mašīnmācīšanās infrastruktūrā.
NVIDIA priekšrocība balstās uz ekosistēmu: draiveri, bibliotēkas, optimizācijas rīki, kopienas pieredze. TorchTPU nepadara TPU ātrākus par GPU. Tas padara tos salīdzināmus, jo izstrādātājs var palaist to pašu PyTorch kodu uz abām platformām bez pārrakstīšanas. TechInformed analītiķi raksta, ka TorchTPU nemaina TPU veiktspēju, bet maina to cenu salīdzināmību. Pircējs, kurš var vienā komandā pārslēgt darba slodzi starp GPU un TPU, sāk izvēlēties pēc cenas par FLOPS.
Pirmais stabils TorchTPU laidiens ar pilnu dokumentāciju vēl nav izziņots. Repozitorijs GitHub ir atvērts ar Apache 2.0 licenci, kas ļauj komerciālu lietošanu bez ierobežojumiem.
Komentāri
Šim rakstam vēl nav komentāru. Esi pirmais, kurš dalās ar savu viedokli.