Google-ն և Meta-ն ներկայացրել են TorchTPU՝ Google TPU-ների համար բնիկ PyTorch բեքենդ

TorchTPU-ն Google-ի և Meta-ի բաց կոդով համագործակցություն է, որը Google TPU-ների համար ներկայացնում է բնիկ PyTorch բեքենդ։ Այն թույլ է տալիս մոդելները գործարկել առանց լրացուցիչ կարգավորումների և տեղափոխել TPU՝ միաժամանակ ապահովելով մրցունակ եզրակացության և ուսուցման արդյունավետություն՝ անմիջական կատարման օպտիմիզացիաների, torch.compile ինտեգրման և սարքային արագացման միջոցով։

Google TPU-ները լայնորեն կիրառվել են մեծ ML մոդելների ուսուցման և տեղակայման համար։ Չնայած TPU-ները պատմականորեն կապված էին TensorFlow-ի և JAX-ի հետ, այս աշխատանքը սկսվել է Google-ում 2025 թվականի կեսին՝ Meta-ի PyTorch թիմի հետ։ Ավելի վաղ PyTorch/XLA ինտեգրումը, ներկայացված 2020 թվականին և հիմնված Lazy Tensors-ի վրա, պահանջում էր խորը փոփոխություններ սկզբնական կոդում, տուժում էր կոմպիլյացիայի լուրջ ուշացումներից, չուներ անմիջական կատարման աջակցություն և հետ էր մնում բաշխված պրիմիտիվներից։

Դիզայնի սկզբունքները և կիրառելիությունը

TorchTPU-ը կառուցվել է որպես սկզբունքային TPU բեքենդ PyTorch-ի համար, այլ ոչ թե որպես JAX-ի հարմարեցում։ Նրա սկզբունքներն են՝ PyTorch-ին պատկանելիությունը (runtime-ում JAX կախվածության բացակայություն, կայուն անմիջական ռեժիմ, Torch Dynamo-ի վրա հիմնված torch.compile, DDP, FSDP և DTensor աջակցություն), էկոհամակարգի անխափան ինտեգրումը և համայնքի ներգրավվածությունը։ Օգտատերերի վրա կենտրոնացումը տեղափոխման պարզություն է՝ model.to("cuda")-ն model.to("tpu")-ով փոխարինելը պետք է թույլ տա մոդելը գործարկել առանց լրացուցիչ կարգավորումների, թեև պիկային արդյունավետության համար կարող է լրացուցիչ կարգավորում պահանջվել։ Նախագիծը նաև ձգտում է սահմանափակ դինամիկության՝ ձևերի փոփոխությունների պատճառով XLA-ի վերակոմպիլյացիայից խուսափելու համար, և ուսումնասիրում է դինամիկ TPU միջուկներ, որոնք կարելի է նախապես կոմպիլյացնել։ TorchTPU wheel փաթեթը տեղադրված TPU մեքենայում PyTorch-ի ներմուծումը պահանջում է միայն չնչին փոփոխություններ PyTorch-ի ինքնաբեռնման մեխանիզմի շնորհիվ։

Անմիջական կատարումը և կոմպիլյացիան

TorchTPU-ն ընդարձակում է PyTorch-ի դիսպետչերը C++ բեքենդով, որը գրանցված է որպես PrivateUse1, և սահմանում է TPU թենսորի պահեստավորումը ստանդարտ at::Tensor մեխանիզմով։ Այն ունի անմիջական կատարման երեք ռեժիմ՝ Strict (լռելյայն) մեկ առ մեկ ասինխրոն ուղարկում է մեկ Aten գործողություն. Debug (TORCH_TPU_LAUNCH_BLOCKING=1) սինխրոն է՝ սահմանների լրացուցիչ ստուգումներով. DeferAndFuse (TORCH_TPU_DEFER_AND_FUSE=1) հետաձգում և միավորում է մի քանի գործողություն, և հաղորդվում է, որ այն 50%-ով կամ ավելի արագ է Strict ռեժիմից՝ միաժամանակ նվազեցնելով HBM հիշողության պիկը։ Մատերիալացման կետում, օրինակ .to("cpu") կանչի ժամանակ, TorchTPU-ն կիրառում է կենսունակության վերլուծություն, կառուցում է քեշի բանալի գործողությունների գրաֆի համար և ուղարկում է քեշավորված ֆայլը կամ ասինխրոն կոմպիլյացնում է։ Repeated Ops Heuristic-ը հայտնաբերում է կրկնվող հաջորդականություններ, օրինակ՝ տրանսֆորմերի բլոկների կույտեր, և մատերիալացումն ավելի վաղ է գործարկում։ Եռաստիճան կոմպիլյացիայի քեշը ընդգրկում է գործընթացի լոկալ հիշողությունը, մեքենայի լոկալ ընդհանուր հիշողությունը և հեռավոր ֆայլային համակարգի մշտական քեշը՝ հոսթերի և վերագործարկումների միջև վերաօգտագործելու համար։

Աջակցվող աշխատանքային ծանրաբեռնվածությունները և ճշգրտությունը

TorchTPU-ն աջակցում է մոտ 450 Aten գործողություն, ներառյալ հիմնական ճարտարապետություններին անհրաժեշտ Core ATen գործողությունների ուղիղ և հակադարձ անցումները, ինչպես նաև հատուկ միջուկներ JAX/Pallas-ի կամ Helion DSL-ի միջոցով։ Այն վավերացվել է Hugging Face-ի ամենաշատ ներբեռնված 5 000 մոդելների վրա՝ առանց լրացուցիչ կարգավորումների, ընդգրկելով LLM-ներ, աուդիո և խոսք, տեսողություն և տեքստային կոդավորիչներ։ DTensor-ի վրա հիմնված շարդինգը, բազմահոստային եզրակացությունը vLLM-ի և SGLang-ի հետ, նուրբ կարգավորումը և ուսուցումը TorchTitan-ով, ինչպես նաև ինտեգրումը Ray-ի և TorchRec-ի հետ նույնպես վավերացվել են։ TPU-ի մատրիցային բազմապատկման միավորները օգտագործում են bfloat16 մուտքեր FP32 կուտակմամբ, իսկ կծկում չկատարող գործողությունները կատարվում են FP32-ով վեկտորային միավորում։ TorchTPU-ն PyTorch-ի ճշգրտության API-ները քարտեզագրում է այս սարքային հնարավորություններին. medium-ը (լռելյայն) օգտագործում է BF16 բազմապատկում FP32 կուտակմամբ. high-ը օգտագործում է 3 անցումով ալգորիթմ՝ մոտ 3 անգամ ավելի հաշվողական ծախսով. highest-ը՝ 6 անցում, մոտ 6 անգամ ավելի ծախսով։ high-ը կամ highest-ը ընտրելը գործարկում է արդյունավետության զգուշացում։

Աղբյուրներ: PyTorch · Google Cloud Blog · Hugging Face Blog