A Google és a Meta TPU-kra vitte a PyTorch nagymodell-tréningjét

2026. október 20.Forrás: Google for Developers

A Google és a Meta elérhetővé tette a TorchTitan TPU-változatát, amely a PyTorch nagy modelljeinek Google Cloud TPU-kon történő betanítását célozza. A megoldás a TorchTPU natív PyTorch-integrációjára épül, és kész recepteket kínál többek között a Llama 3, a Qwen 3 és a FLUX modellcsaládokhoz.

A lényeg röviden

A TorchTitan TPU Google Cloud TPU-kon futó PyTorch-modelltréninghez készült.
A TorchTPU natív PyTorch-backendet ad a TPU-khoz a torch.device("tpu") használatával.
Kész receptek érhetők el többek között Llama 3, Qwen 3 és FLUX modellekhez.
A Graph Trainer a teljes tanítási lépést egyetlen XLA-programként rögzítheti.
A Google szerint a Qwen 3 235B 29 százalék feletti MFU-t, a Qwen 3 8B pedig 51 és 52 százalék közötti MFU-t ért el bizonyos útvonalakon.

Natív TPU-támogatás a PyTorchban

A Google Developers Blog október 20-i bejegyzése szerint a Google Cloud TPU-k a TorchTPU révén natív PyTorch-backendet kaptak. A fejlesztők a torch.device("tpu") eszközmegadásával használhatják a TPU-kat, így egy meglévő tanítási szkript eszközkezelésének módosítása általában csak kisebb változtatást igényel.

A TorchTitan a Meta PyTorch-alapú, nagy léptékű generatívmodell-tréninghez készült referenciaimplementációja. A Google és a Meta által közösen fejlesztett TPU-változat, a TorchTitan TPU, a Meta megoldásának olyan elágazása, amely Google Cloud TPU-kon teszi lehetővé a modellek tanítását. A kód a google-pytorch/torchtitan GitHub-tárhelyen érhető el.

A Google szerint korábban a PyTorch-modellek TPU-ra vitele vagy teljes kódátírást igényelt JAX-ra, vagy a régebbi torch_xla használatával járt. A TorchTPU ezt a folyamatot közvetlenül a PyTorch felületéhez kapcsolja.

Kész receptek nagy modellekhez

A TorchTitan TPU célja, hogy a fejlesztőcsapatoknak ne kelljen saját TPU-s particionálási stratégiákat vagy alacsony szintű XLA-kommunikációs megoldásokat készíteniük. A keretrendszer hivatalos recepteket tartalmaz sűrű modellek előtanításához, szakértőkeverék, vagyis MoE-modellekhez, valamint diffúziós transzformerekhez.

A felsorolt példák között szerepel a Llama 3 70B és 405B, a Qwen 3 8B és 32B, a Qwen 3 235B MoE, valamint a FLUX.1 12B. A TorchTitan a tanítás utáni folyamatokban is használható, például a VeRL GRPO- és megerősítéses tanulási folyamataiban. Ilyenkor a paraméterfrissítést végzi, miközben a súlyokat olyan nagy áteresztőképességű kiszolgáló motorokkal szinkronizálja, mint a vLLM-TorchTPU.

A modell definíciója és a modell particionálása elkülönül egymástól. A modellek tiszta, egyetlen eszközre írt PyTorch-kódon alapulnak, a párhuzamosítás pedig keretrendszerszinten, deklaratív particionálási tervekkel vagy TPU-specifikus bővítményekkel adható hozzá.

TPU-ra optimalizált végrehajtás és fordítás

A TorchTitan konfigurációs fájából a modell felépítése előtt egyes blokkok lecserélhetők. A TPU-s receptek így többek között Splash Attentiont használnak, amely a Google leírása szerint a maszkok kezeléséhez nem materializálja a teljes figyelmi mátrixot. A végső projekció és a keresztentrópia-veszteség összevont Pallas-kernellel váltható ki, amely nem írja ki a teljes logit-tenzort a nagy sávszélességű memóriába. Az MoE-szakértők mátrixműveleteihez csoportosított Pallas-kódok állnak rendelkezésre.

A TorchTPU Pallas-kernelfunkciókat PyTorch egyedi műveletekké alakít, ehhez a torch.tpu.pallas.jax_op dekorátort használja. A cserék konfigurációs kapcsolókkal aktiválhatók, és tesztekkel vethetők össze az eredeti PyTorch-megoldással.

A Google külön SimpleFSDP-megoldást is bemutat. Ez a paraméterek particionálását a fordító által látható adatfolyamként kezeli, ezért a paraméterek összegyűjtése és a gradiens visszaosztása a nyomkövetett gráf része lehet. A leírás szerint ez lehetővé teszi, hogy az XLA a kommunikációt a számításokkal együtt ütemezze. A FLUX esetében a megoldás a JAX MaxDiffusion teljesítményének 99 százalékos paritását érte el.

Egyetlen gráfba foglalható a teljes tanítási lépés

A TorchTitan Graph Trainer eszköze a teljes tanítási lépést, a közvetlen továbbítást, a veszteségszámítást és a visszaterjesztést egyetlen FX-gráfban rögzíti a PyTorch make_fx nyomkövetőjével. A lépés ezután egyetlen XLA-programként hajtható végre. A megközelítés aktivációs memóriakezelési szabályokat is alkalmazhat, amelyek eldöntik, mit kell menteni, újraszámolni vagy gazdagépi memóriába kihelyezni.

A Google által közölt eredmények szerint a Qwen 3 235B ezen az útvonalon 29 százalék feletti MFU-t ér el, míg a Qwen 3 8B esetében a funkcionális AdamW-frissítést is a tanítási gráfba foglaló útvonal 51 és 52 százalék közötti MFU-t ér el. A fejlesztők számára a legtöbb optimalizáció konfigurációs kapcsolóval aktiválható, így egy meglévő recept finomhangolása sok esetben új kód írása nélkül elvégezhető.

Ez a felhasználóknak elsősorban azt jelenti, hogy a Google Cloud TPU-kra szánt PyTorch-modellekhez kész, nagy modellekre szabott tréningmegoldások állnak rendelkezésre. A TorchTPU a kompatibilitást, a TorchTitan TPU pedig az elosztott, teljesítményre hangolt tanítási recepteket adja hozzá ehhez a munkafolyamathoz.