- Google'i TorchTPU ja PyTorch/XLA muudavad TPU-d PyTorchi natiivseks ja suure jõudlusega taustaprogrammiks ilma JAX-stiilis vaimset mudelit sundimata.
- TPU arhitektuur, XLA kompileerimine ja StableHLO võimaldavad tõhusat tihedat arvutust ja kollektiive massilisel skaalal, eriti hajutatud treeningu puhul.
- Uued innukad režiimid, piiratud dünaamilisus ja ökosüsteemi tööriistad nagu easy-torch-tpu vähendavad hõõrdumist GPU-keskse PyTorch-koodi migreerimisel TPU-klastritesse.
- Cloud TPU, GKE ja Vertex AI pakuvad infrastruktuuri, et käitada TPU-del kõike alates uurimistöö ulatusest kuni pod-skaala PyTorchi töökoormusteni.
PyTorchi käitamine Google TPU-des ei ole enam nišš, eksperimentaalne tee, mis on reserveeritud käputäiele ekspertideleGoogle'i uue vahel TorchTPU virn, lahingutes testitud PyTorch/XLA projekt ning kasvav tööriistade ja raamistike ökosüsteem, mis treenib ja serveerib mudeleid TPU-del, on kiiresti muutumas sama loomulikuks kui NVIDIA GPU-dega töötamine. Suur muutus seisneb selles, et nüüd saab pürgida samaaegselt suure jõudluse, tohutu ulatuse ja palju sujuvama arendajakogemuse poole.
See artikkel süveneb sellesse, kuidas PyTorch tänapäeval TPU-sid kasutab ja kuhu see virn liigub.Me lahti pakime TorchTPU arhitektuuri, erinevused traditsioonilise PyTorch/XLA-ga võrreldes, kuidas hajutatud treening, kompileerimine ja riistvara eripärad toimivad ning mida see praktikas tähendab GPU-kesksete PyTorchi töövoogude migreerimisel. Kui elate õigusteaduslike matemaatiliste meetodite, difusiooni või suuremahuliste soovitussüsteemide maailmas, siis allolevad üksikasjad on just selline madala taseme reaalsus, mis otsustab, kas teie TPU töötab lennates või roomates.
Miks PyTorch TPU-del praegu oluline on?
Kaasaegsed tehisintellekti töökoormused on lihtsast „üks masin, vähe graafikaprotsessoreid” ajastust välja kasvanud.Tipptasemel mudelid ulatuvad nüüd klastritesse, mis sisaldavad kümneid tuhandeid kiirendeid, sundides tarkvara hakkama saama äärmise ulatusega, usaldusväärse hajutatud täitmise ja kaasaskantava jõudlusega erinevate kiipide ja müüjate vahel. AI infrastruktuur.
Google'i tensorprotsessorid (TPU-d) on selle piiri keskmes.Nad toetavad sisemisi süsteeme nagu Gemini ja Veo, aga ka suurt osa Google Cloudi klientide koolitus- ja järelduskoormustest. Ajalooliselt olid TPU-d tihedalt seotud JAX-i ja TensorFlow'ga, kuid laiem ökosüsteem on standardiseerinud tugevalt PyTorchi, mis on tekitanud valuliku lõhe: GPU-d tähendasid „PyTorch + CUDA”, TPU-d aga „JAX + XLA”.
Google'i vastus on täisgaasiga pingutus, et TPU-d tunneksid end esmaklassilise PyTorchi sihtmärginaTorchTPU eesmärk on pakkuda teile natiivset ja innukat PyTorchi semantikat tipptasemel jõudlusega, samas kui PyTorch/XLA jääb võimsaks ja laiskalt kompileeritavaks lahenduseks, mis on juba laialdaselt tootmiskeskkonnas kasutusele võetud. Nende pakettide ümber muudavad Cloud TPU, GKE, Vertex AI ja kogukonnaraamistikud nagu easy-torch-tpu TPU klastrid lihtsaks ja skriptitavaks infrastruktuuriks kõigele alates 1B kuni 70B+ parameetriga mudelitest.
TPU riistvara sees: enamat kui lihtsalt kiirem kiip
TPU-süsteem on põhimõtteliselt tihedalt integreeritud kiipide, hostide ja ühenduste struktuur., mitte ainult üks kiirenduskaart. Selle paigutuse mõistmine on oluline, et mõista TorchTPU disaini ja seda, miks selle kompilaatori valikud erinevad puhastest GPU-pakkidest.
Iga TPU host ühendub mitme TPU kiibiga kiipidevahelise ühenduse (ICI) kaudu.ICI moodustab suure ribalaiusega 2D- või 3D-toruse topoloogia, mis võimaldab suuremahulistel podidel käituda ühe loogilise kiirendi rollis. Gradientide edastamise asemel traditsiooniliste võrgukihtide kaudu liiguvad kollektiivid otse sellel torusel, muutes skaleerimise palju tõhusamaks, kui teie tarkvara teab, kuidas neid kollektiive õigesti väljendada.
TPU kiibi sees on arvutusvõimsus jagatud TensorCore'ide ja SparseCore'ide vahel.TensorCore'id on spetsiaalsed ühelõimelised mootorid, mis paistavad silma tiheda maatriksi matemaatikaga – täpselt sellega, mis annab jõudu transformaatoritele, CNN-idele ja enamikule standardsetele süvaõppe kihtidele. SparseCore'id on loodud töökoormuste jaoks, millel on ebaregulaarsed mälule juurdepääsu mustrid, näiteks manustamine, kogumine/hajutamine ja koormatud kollektiivsed toimingud.
See arhitektuur on süvaõppe jaoks fantastiline, kuid selle toiteallika osas on see valiv.Näiteks paljud transformaatorite implementatsioonid panevad kõvakoodi abil paika pea mõõtmed 64. Praegused TPU põlvkonnad saavutavad oma optimaalseima tulemuse 128–256 vahel, mis tähendab, et pea mõõtmete kahekordistamine võib oluliselt parandada maatriksite korrutamise efektiivsust ja TensorCore'i kasutamist. Kaasaskantavus ei kustuta neid riistvaralisi tegelikkusi, see lihtsalt lihtsustab nende saavutamist.
PyTorch/XLA-st TorchTPU-ni: kaks teineteist täiendavat viisi PyTorchi käitamiseks TPU-del
PyTorch saab juba täna TPU-del PyTorch/XLA (torch_xla) kaudu töötada., mis esitleb TPU-sid standardsete PyTorch-seadmetena ja kompileerib kapoti all laisaid XLA-graafe. Paljud teadlased on aga leidnud, et kuigi nende koodi muudatused on paberil väikesed, võib käitumise erinevus võrreldes GPU-innuka täitmisega tunduda häiriv.
TorchTPU on Google'i uus natiivne PyTorchi taustprogramm, mis on loodud tunduma nagu „päris” PyTorch, mitte ümbrisSelle asemel, et sundida PyTorchi JAX-i-laadsele mudelile, kus kõikjal on laisad tensorid, tugineb TorchTPU PyTorchi kiirele täitmisele ja kaasaegsetele kompileerimise API-dele, näiteks torch.compile. See kasutab Privaatne kasutamine1 seadme mehhanism PyTorchis, seega teie vaatenurgast töötate lihtsalt tavalistega taskulamp.Tensor objektid, mis juhtuvad TPU-l elama.
Kahe lähenemisviisi peamine erinevus on teostusstiilPyTorch/XLA kasutab vaikimisi laisat teostust: operatsioonid loovad graafi, mis seejärel käivitab XLA kompileerimise, kui jõuate sünkroonimistõkkeni, näiteks treeningtsükli sammuni. TorchTPU seevastu on üles ehitatud kui „Eager First“, millel on lisarežiimid, mis järk-järgult ühendavad operatsioone ja edastavad optimeeritud alamgraafe XLA-le ilma, et peaksite loobuma standardsest PyTorchi mentaalsest mudelist.
Pilvepõhine TPU, GKE ja Vertex AI: infrastruktuuri selgroog
Iga valitud PyTorch-on-TPU virna all on Cloud TPU platvorm., mis pakub kohandatud ASIC-e skaleeritavate pilveressurssidena, mis on häälestatud nii treenimiseks kui ka järelduste tegemiseks. Neid kiirendeid kasutatakse väga erinevate töökoormuste jaoks: vestlusagendid, koodi genereerimine, pildi- ja meediamudelid, kõne, soovitussüsteemid ja isikupärastamismootorid.
Pilvepõhised TPU-d on tihedalt integreeritud Google Kubernetes Engine'iga (GKE), nii et saate ajastada suuremahulisi PyTorchi töid standardsete Kubernetes'i primitiivide abil. Dünaamiline töökoormuse ajastaja võimaldab teil taotleda kogu vajalike kiirendite laevastikku korraga, tagades, et tuhanded TPU-kiibid on koos võrgus, et treenida või teenindada mudelit ilma käsitsi orkestreerimiseta.
Meeskondade jaoks, kes soovivad kõige lihtsamat sissejuhatust, võtab Vertex AI ära suurema osa klastri haldamisest.Saate sihtida TPU-sid hallatud koolitus- ja teenindusvoogude kaudu, sealhulgas siis, kui kasutate PyTorchi-põhised mudelidGoogle Cloud positsioneerib seda paindlikkust – TPU-sid või GPU-sid, hallatud või isetehtud Kubernetes'i – otsese vastusena nii ettevõtete kui ka uurimislaborite plahvatuslikult kasvavale nõudlusele tehisintellekti infrastruktuuri järele.
TorchTPU põhifilosoofia: „PyTorchi kodakondsus”
TorchTPU keskne disainieesmärk on otsekohene: see peaks tunduma nagu PyTorch, mitte nagu võõras raamistikKui sa juba tead, kuidas CUDA GPU-del mudelit treenida, peaksid sa saama sama treeningskripti TPU-desse portida minimaalsete koodimuudatustega ja ilma oma mentaalset mudelit ümber kirjutamata.
Praktikas tundub ideaalne ränne peaaegu koomiliselt lihtneKuhu sa tavaliselt kirjutaksid seade = torch.seade('cuda'), saate selle asemel TorchTPU moodulist TPU-seadme – põhimõtteliselt midagi sellist seade = tpu.get_device()– ja helista mudel.to(seade) täpselt nagu graafikakaardil. Sinu edasisööt, optimeerija loogika ja Hugging Face'i mudelite kutsumise viis jäävad samaks.
Varasemad TPU integratsioonid sundisid PyTorchi sageli JAX-i jäljendamaNad toetusid suuresti laisatele tensoritele ja sundisid sind staatiliste graafide mõtlemisele. See murdis ühe PyTorchi suurimatest tugevustest: sa ei saanud lihtsalt edastuskäigu keskele printimist sisestada, et kujundeid või väärtusi kontrollida. TorchTPU lükkab selle kompromissi tagasi. See hoiab innukat käitumist baasjoonena ja ehitab jõudlust selle ümber, selle asemel, et paluda sul sellest loobuda.
See „PyTorchi kodakondsuse” põhimõte laieneb ka veakäsitluseleXLA pinu sügavusse peidetud krüptiliste 500-realiste C++ pinu jälgede asemel on eesmärk leida puhtaid Pythoni tagasijälgi, mis osutavad otse teie treeningtsükli või mudeli definitsiooni vigasele reale. Kui žongleerite mitme miljardi parameetriga mudelite ja tuhandete TPU-dega, on see elukvaliteedi paranemine see, mis eristab pärastlõunast parandust sihitutest silumispäevadest.
TorchTPU innukad režiimid: silumine, range ja ühendatud
Suurte sulandgraafikute jaoks loodud riistvaral natiivse ja innukalt toimiva kogemuse pakkumine pole triviaalne.TorchTPU lahendab selle, pakkudes mitmeid innukalt toimivaid režiime, mida toetab jagatud kompileerimis- ja täitmistorustik, nii et saate sujuvalt liikuda režiimilt „pane see tööle“ režiimile „tee see kiiresti“.
Silumine innukas on kõige aeglasem, kuid samas kõige läbipaistvam režiim. See saadab üks operatsioon korraga TPU-le ja sünkroniseerub protsessoriga pärast iga toimingut. Jõudlust ohverdatakse tahtlikult, et saaksite hõlpsalt jälgida NaN-e, kuju mittevastavusi või mälu otsasaamise vigu kohese tagasiside ja selgete pinujälgede abil.
Range innukas säilitab selle üheoperatsioonilise lähetuse semantika, aga käivitab asünkroonseltTPU ja CPU saavad töötada paralleelselt, kuni kasutajakood jõuab sünkroniseerimispunktini, pakkudes kogemust, mis on palju lähedasem tavalisele GPU-toega innukale PyTorchile, kuid ilma suurte graafikute kompileerimise nõueteta.
Fused Eager on koht, kus asjad lähevad esituse seisukohast tõeliselt huvitavaksTorchTPU jälgib teie teostatavate toimingute voogu ja ühendab need automaatselt suuremateks ja tihedamateks arvutusplokkideks enne XLA kaudu TPU-le saatmist. See dünaamiline ühendamisetapp suurendab oluliselt TensorCore'i kasutamist ja vähendab mälu ribalaiuse üldkulu, andes rutiinselt tulemuseks 50–100%+ kiirendused võrreldes Strict Eageriga ilma mudelikoodi muudatusteta.
Kõigil kolmel innukal režiimil on ühine kompileerimisvahemälu mis võib asuda ühel hostil või olla püsiv mitmel hostil hajutatud keskkonnas. Aja jooksul, kui teie treeningtsükkel stabiliseerub ja süsteem näeb samu mustreid, langevad kompileerimise kulud ning te kulutate rohkem aega tensorite analüüsimisele, mitte käivitatavate failide loomisele.
Staatiline kompileerimine: torch.compile, XLA ja StableHLO
Kui vajate TPU-del absoluutset tippjõudlust, lülitub TorchTPU otse kaasaegsesse PyTorchi kompileerimistorustikku.Saate mudeleid või funktsioone mähkida torch.compile(), mis jäädvustab Torch Dynamo abil FX-graafiku, seejärel möödub tavapärasest TorchInductori taustaprogrammist ja annab juhtimise hoopis XLA-le.
XLA valimine peamiseks serveripoolseks süsteemiks on teadlik otsus, mis on juurdunud TPU reaalsusesse.XLA-d on aastatepikkuse juurutamise käigus TPU pod'ides karastatud ning see mõistab sügavalt tiheda matemaatika ja kollektiivse suhtluse kokkupuutepunkti ICI toruses. TorchTPU kaardistab PyTorchi operaatorid otse Stabiilne HLO, OpenXLA poolt mõistetav tensor-infrapuna, laseb seejärel XLA langetavatel käikudel genereerida optimeeritud TPU binaarfaile, taaskasutades võimaluse korral samu käitusteid kui innukad režiimid.
Kohandatud operaatorite laiendatavus ei ole teisejärguline mõteTorchTPU toetab Pallases ja JAX-is määratletud kohandatud kerneleid: JAX-funktsiooni kaunistades millegi sarnasega @torch_tpu.pallas.custom_jax_kernel, saate kompileerimisteesse süstida madala taseme riistvaraliselt häälestatud koodi, kaotamata globaalse optimeerija eeliseid. Samuti on käimas töö täiendavate DSL-ide, näiteks Helioni, toetamiseks, et kerneli loomist veelgi paindlikumaks muuta.
Hajutatud PyTorch TPU-del: DDP, FSDP, DTensor ja MPMD
Massiivsed mudelid ei treeni ühe kiirendi abil ja TorchTPU on ehitatud seda reaalsust silmas pidades.See integreerub otse PyTorchi standardsete hajutatud API-dega, sealhulgas Hajutatud Andmed Paralleelselt (DDP), FSDPv2ja DTensorja on valideeritud kolmandate osapoolte teekidega, mis tuginevad nendele abstraktsioonidele.
Üks PyTorch/XLA suuri ajaloolisi valupunkte oli selle range SPMD (üksikprogramm, mitu andmemahtu) kallutatus.Paljudel reaalsetel PyTorchi treeningskriptidel on astmete vahel väikesed erinevused – 0. aste võib käsitleda logimist, kontrollpunkte või mõõdikuid, samas kui teised astmed teevad puhast arvutust. XLA globaalse graafi vaate puhul oli selline käitumine kohmakas ja sundis arendajaid sageli koodi ümber kirjutama, et erinevusi vältida.
TorchTPU hõlmab selgesõnaliselt MPMD (mitme programmi, mitme andmestiku) stsenaariumeSee isoleerib ja piirab hoolikalt kommunikatsiooniprimitiivid, et erinev käitumine ei rikuks korrektsust ega vähendaks jõudlust. Võimaluse korral võimaldab see XLA-l endiselt näha hajutatud arvutuse globaalset pilti, et katta suhtlus arvutusega, kuid see ei sunni teid enam ebareaalselt puhta SPMD-stiili poole.
Eriti oluline on see, kuidas see haakub olemasolevate PyTorchi hajutatud paradigmadega.Raamistikud nagu FSDP, DTensor ja ökosüsteemi tööriistad nagu TorchTitan tuginevad sellele. Protsessigrupp API kollektiividele nagu all-reduce, all-gather ja broadcast. Graafikaprotsessorite puhul suunatakse need kõned tavaliselt NCCL-i. TorchTPU pealt kuulab need kollektiivid ProcessGroupi kihil ja alandab need StableHLO kollektiivoperatsioonideks, mida TPU riistvara ja ICI torus natiivselt täidavad. FSDP või DTensori vaatenurgast pole midagi muutunud – nad näevad lihtsalt teistsugust tausta.
PyTorch/XLA: laisk teostus, sünkroniseerimispunktid ja praktilised näpunäited
Kuigi TorchTPU on pikaajaline ja täielikult natiivne lahendus, on PyTorch/XLA endiselt peamine tööriist PyTorchi käitamiseks TPU-del.Kui olete CUDA innuka teostusega harjunud, siis PyTorch/XLA suurim kontseptuaalne nihe on see, et tensorid on laiskToimingud salvestavad graafi; tegelik täitmine ja kompileerimine toimuvad siis, kui te otseselt või kaudselt sünkroonite.
Sünkroniseerimispunktid on kohad, kus PyTorch/XLA annab loodud graafi XLA-le kompileerimiseks ja käivitamiseks üle.Tüüpilisteks takistusteks on sellised kõned nagu torch_xla.sync() või kõrgema taseme utiliidid, näiteks xm.optimizer_step(optimeerija), mis nii astmeliselt optimeerivad kui ka sünkroniseerivad gradiente seadmete vahel, kui olete hajutatud seadistuses.
Sellel laisal mudelil on suured tagajärjed jõudluseleAntud graafi (või uute sisendkujudega graafi) esmakordsel käivitamisel maksate kompileerimistasu, kuid järgnevad iteratsioonid toimuvad palju kiiremini, kui struktuur jääb stabiilseks. Seetõttu on kuju stabiilsus – fikseeritud jadade pikkused, järjepidevad partiide suurused – PyTorch/XLA töökoormuste jaoks nii oluline ja miks... sisendite täitmine fikseeritud suurustega on nii levinud muster.
Mitmeprotsessiline treening PyTorchis/XLA-s kasutab omaenda mugavustööriistuTavaliselt mässid sa oma põhitreeningu funktsiooni (näiteks _mp_mnist_fn) ja käivitage see kõigis seadmetes rakendusega torch_xla.launchAndmete laadimist hallatakse kaudu. torch_xla.distributed.parallel_loader.MpDeviceLoader, mis võtab standardse PyTorch DataLoader'i ja tagab, et iga protsess näeb unikaalset andmekildu, samal ajal partiisid vastavasse TPU-seadmesse eeltellides.
Andmete laadimine, hajutatud täitmine ja AMP TPU-del
Tõhusad sisendkanalid on TPU-de puhul sama olulised kui GPU-de puhul.PyTorch/XLA-l MpDeviceLoader kattub hostipoolse andmete laadimise ja seadmepoolse täitmisega, edastades partiisid otse TPU-le ja aidates teil vältida pikki jõudeolekuperioode, kui kiirendi ootab uusi andmeid.
Hajutatud treeningu puhul teeb xm.optimizer_step(optimizer) enamat kui lihtsalt optimeerija samm.See teostab gradient-kõik-vähenduseid kõigis seadmetes, arvutab nende keskmise, rakendab kaaluuuendusi ja tegeleb vajaliku sünkroniseerimisega, seega ei ole tavaliselt vaja igas iteratsioonis eraldi selgesõnalist sünkroonimiskõnet. Logimise abilised, näiteks xm.is_master_ordinal(local=False) Veenduge, et ainult üks protsess tegeleks mõõdikute ja kontrollpunktidega, et vältida dubleerimist.
Automaatne segatäpsus (AMP) näeb TPU-del ja GPU-del veidi teistsugune välja.TPU-d toetavad natiivselt bfloat16 (BF16), mis pakub palju suuremat eksponentide vahemikku kui float16 ja tavaliselt ei vaja stabiilsuse tagamiseks selget kadude skaleerimist. PyTorch/XLA laiendab PyTorch AMP-i, et see vajadusel automaatselt BF16 ja FP32 vahel kaardistaks, muutes segatud täpsusega treenimise TPU-del nii lihtsaks kui ka töökindlaks.
Mudelite salvestamisel on ka TPU-spetsiifiline parim tavaKuigi saate helistada taskulamp.save seadme tensoritest on üldiselt soovitatav vii olekuteated enne serialiseerimist protsessorisse PyTorch/XLA kasutamisel, mis muudab nende laadimise mitte-TPU riistvaral, näiteks tavalistel GPU-masinatel, lihtsamaks.
Lihtsad taskulamp-TPU ja reaalse TPU treeningraamistikud
Lisaks ametlikele raamistikele ehitab kogukond kõrgema taseme raamistikke, et muuta TPU-de omaksvõtmine lihtsamaks.. Üks näide on aklein4/easy-torch-tpu, kerge treeningraamistik, mis on loodud spetsiaalselt PyTorch/XLA töövoogude lihtsustamiseks Google Cloud TPU klastrites.
Easy-torch-tpu positsioneerib end lihtsama ja paindlikuma alternatiivina suurtele ja jäikadele koodibaasidele nagu Hypercomputer/torchprime.Selle disainiprioriteedid on selged: lihtne seadistamine, otsekohene kohandatavus ja sujuv integreerimine gcloud ssh-põhised klastri töövood. See on sihilikult suunatud „akadeemilise ulatusega“ katsetele – mudelitele parameetrite vahemikus 1–10B umbes 32–64 TPU kiibil.
Laiendatavust käsitletakse alamklasside ja konfiguratsioonifailide kauduUute alamklasside lisamisega saate ühendada oma arhitektuure, treeningtsükleid, optimeerijaid, andmelaadureid ja isegi kohandatud killustamis- ja taasmaterialiseerimisstrateegiaid. See võimaldab teil vabalt katsetada, taaskasutades raamistiku hajutatud ja logimisraamistikku.
Raamistik integreerub tihedalt peamiste ökosüsteemi tööriistadegaKaalude ja eelarvamuste tugi muudab katse jälgimise lihtsaks, samas kui Hugging Face'i integratsioon lihtsustab andmestike laadimist, eeltreenitud kontrollpunktide hankimist ja mudelite salvestamist, mida saab hiljem käivitada standardsel GPU-põhisel PyTorchil. Repositoorium sisaldab installidokumentatsiooni, alustusnäiteid ja areneb aktiivselt kogukonna tagasiside põhjal.
Piirangud, silumine ja jõudluse lõksud
Isegi kõigi nende täiustuste korral pole PyTorchi käitamine TPU-del veel täiesti hõõrdevabaSuurte mudelite või dünaamiliste töökoormuste käitlemisel säästab palju aega mõistmine, kus asjad võivad valesti minna.
Graafikute ümberkompileerimine on endiselt üks suurimaid varjatud jõudluse tapjaidIga kord, kui teie arvutusgraafik või sisendkujud sünkroonimispunktide vahel muutuvad, võib XLA olla vaja uuesti kompileerida, mis tekitab märgatavaid pause. See on eriti levinud muutuva pikkusega jadade või adaptiivsete partiide suuruste puhul, mis on tavalised keele modelleerimise ja genereerimise töökoormustes.
Toetamata või osaliselt toetatud operaatorid võivad jõudlust märkamatult vähendadaKuigi PyTorch/XLA ja TorchTPU eesmärk on laia operaatori katvus, ei pruugi mõnel Aten-operatsioonil veel olla natiivseid XLA-langetusi. Sellistel juhtudel võib täitmine langeda tagasi protsessori peale, mis on tehniliselt korrektne, kuid võib olla suurusjärkude võrra aeglasem. Sisseehitatud silumisutiliidid ja mõõdikud (näiteks torch_xla.debug.metrics) aitavad teil tuvastada protsessori varuprobleeme või ootamatuid ümberkompileerimisi.
Klassikalised GPU profileerimistööriistad nagu Nsight ja nvprof ei näe TPU kernelite sisemustSelle asemel loodate kitsaskohtade mõistmiseks XLA-spetsiifilistele profileerimiskonksudele, TPU käitusaja mõõdikutele ja kõrgema taseme logimisele. Paljud meeskonnad leiavad, et kui nad on parimate tavade (staatilised kujundid, hoolikas andmete laadimine, rekompileerimise jälgimine) omaks võtnud, jõuavad nad kiiresti prognoositava jõudluseni.
Google'i kompilaatori tegevuskava on otseselt suunatud nendele valupunktideleXLA-s täiustatud piiratud dünaamilisuse kallal töötamise eesmärk on võimaldada mudelitel käsitleda erineva pikkusega jadasid ja partiisid ilma uusi kompileerimisi käivitamata. Kasvav eelkompileeritud TPU tuumade teek püüab vähendada külmkäivituse latentsust uute graafikute esimesel iteratsioonil.
Tegevuskava ja ökosüsteem: hõõrdevaba PyTorchi suunas TPU-del
Tulevikku vaadates on Google'i TorchTPU tegevuskava ambitsioonikas ja tihedalt kooskõlas laiema PyTorchi ökosüsteemiga.Plaanis on avalik GitHubi repositoorium koos ulatusliku dokumentatsiooni, arhitektuuriõpetuste ja reprodutseeritavate näidetega, mis hõlmavad nii koolitus- kui ka teenindusstsenaariume.
Integratsioon PyTorchi Helion DSL-iga on silmapiiril, mis peaks laiendama arendajate valikuid kohandatud TPU tuumade kirjutamiseks ilma XLA või riistvarapõhise koodi sügavaimatesse kihtidesse sukeldumata. Natiivne, esmaklassiline tugi dünaamilistele kujunditele läbi torch.compile on samuti prioriteet, mis peegeldab tänapäevaste järjestuspõhiste mudelite tegelikkust.
Mitme järjekorra tugi on veel üks oluline fookusvaldkondPaljud PyTorchi koodibaasid tuginevad suuresti asünkroonsetele teostusmustritele ja lahutatud mälu-/arvutusvoogudele. Nende idioomide selge ja ilma suuremate ümberarvutusteta TPU-dele kaardistamine vähendab oluliselt suurte ja küpsete projektide migratsioonihõõrdeid.
Sügavad ökosüsteemide integratsioonid on juba käimasKäimas on jõupingutused tugeva skaleerimise valideerimiseks täieliku TPU Pod-i suuruseni ja ühendumiseks suuremate PyTorch-põhiste süsteemidega nagu vLLM ja TorchTitan. Samal ajal teeb Google tihedat koostööd Meta ja PyTorchi kogukonnaga ning uurib TorchTPU põhiosade avatud lähtekoodiga muutmist, et kiirendada kasutuselevõttu ja läbipaistvust.
Kõik see toimub suurema ärikeskkonna taustal, kus TPU võimsus suureneb dramaatiliselt.Google Cloud sõlmib uusi mitme miljardi dollari suuruseid tehisintellekti taristulepinguid, Anthropic plaanib juurdepääsu kuni miljonile TPU-le (võimsusega umbes gigavatt) ja Google müüb TPU-sid isegi otse kohapealsetele andmekeskustele. Aeg, mil TPU-d olid nišš, ainult Google'ile kuuluv sisemine ressurss, on ammu möödas.
Kõike kokku võttes liigub PyTorch-on-TPU lugu märkimisväärselt kiiresti „veidrast kõrvalteest“ „standardvalikuks“.TorchTPU loomuliku innukalt toimiva kogemuse, PyTorch/XLA lahingutesti läbinud laisa teostuse, selliste raamistike nagu easy-torch-tpu ja neid ümbritseva rikkaliku Cloud TPU infrastruktuuri abil saate nüüd võtta peavoolu PyTorchi mudeleid – sageli vaid seadme stringi muutmisega – ja neid tõhusalt käitada mõnel suurimal saadaoleval tehisintellekti superarvutil. Mida rohkem pinu koondub tuttavatele PyTorchi idioomidele uute vaimsete mudelite pealesurumise asemel, seda realistlikumaks muutub riistvaravaliku käsitlemine pigem rakenduse detailina kui põhimõttelise disainipiiranguna.
