A PyTorch natív Trainium-támogatást kap az AWS TorchNeuron révén

A TorchNeuron a PyTorch szabványos eszközére épül

A PyTorch október 19-én jelentette be a TorchNeuront, amely nyílt forráskódú, natív PyTorch-hátteret biztosít az AWS Trainiumhoz. Az AWS Trainium olyan AI-chipek portfóliója, amelyet a vállalat a következő generációs, felhős léptékű AI-modellek tanítására és futtatására tervezett.

A megoldás közvetlenül a PyTorch hardverfüggetlen bővítési pontjaira épül, köztük a PrivateUse1 eszköz-háttérre és a torch.accelerator API-kra. Ennek köszönhetően a fejlesztők a már használt modellkódot, könyvtárakat és eszközöket futtathatják Trainiumon. A TorchNeuron a PyTorch saját API-jain keresztül kezeli többek között a torch.nn, a tenzorokat, az automatikus pontosságkezelést, a megosztott számításokat és a teljesítményprofilozást.

A projekt nyílt forráskódú kiadása Apache 2.0 licenc alatt érkezik a github.com/aws-neuron/torch-neuronx címen, és október 20-án nyilvános bétaként válik elérhetővé.

Kutatáshoz és éles futtatáshoz is kínál módot

A TorchNeuron támogatja a PyTorch két fő végrehajtási módját. Az eager végrehajtás azonnal futtatja a műveleteket, ezért a Python vezérlési elemei, a kiíratások és a töréspontok is használhatók. Ezt a módot a PyTorch a kutatáshoz és hibakereséshez szánja. A fejlesztők számára egy kísérleti adaptív eager mód is készül, amely a műveletindítás költségét csökkentené, és kódmódosítás nélkül közelítené a lefordított gráf teljesítményét.

Éles, nagy teljesítményű modellekhez a torch.compile használható. A Neuron-háttér a modellt gráfba rendezi, majd a Neuron fordító optimalizálja azt. A teljes modell vagy csak egy része is fordítható, az eager móddal együtt. A FlexAttention natív támogatást kapott, az inference-verem pedig ugyanerre a torch.compile-integrációra épül.

A Trainium-eszközre váltás egyszerű: egy modell a .to('neuron') hívással helyezhető át. A PyTorch eszközfüggetlen API-jaira épülő modern kód módosítás nélkül futtatható. Fordított modellnél a fejlesztő egyszer megadhatja a neuron hátteret, vagy közvetlenül a torch.compile hívásban is beállíthatja.

Elosztott tanítás, egyedi kernelek és profilozás

A TorchNeuron a szöveges, multimodális, képi és videós modellekhez, valamint állapottér- és világmodellekhez használt aten-műveletek széles körét támogatja. A támogatott adattípusok között szerepel a BF16, az FP32 és az FP16, továbbá az egész és logikai típusok. A TorchAO és az FP8 jelenleg privát bétában érhető el, általános hozzáférésük a következő kiadásban várható.

Az elosztott tanítás és következtetés a standard PyTorch API-kra épül. A felsorolt támogatott megoldások között van a DDP, az FSDP2, a SimpleFSDP, a tenzor- és csővezeték-párhuzamosítás, valamint ezek kombinációi. A DTensor és a DeviceMesh a particionálást és a többdimenziós folyamatcsoportokat kezeli. A kollektív műveletek a Trainium dedikált kollektív magjain, csomópontok között pedig az Elastic Fabric Adapteren futnak.

A Neuron Kernel Interface, vagyis az NKI lehetővé teszi egyedi műveletek fejlesztését és futtatását eager, illetve torch.compile módban is. A fejlesztők így alacsony szintű hardvervezérlést igénylő ötleteket próbálhatnak ki, vagy teljesítménykritikus NKI-kernelekkel válthatnak ki PyTorch-műveletsorozatokat. A TorchNeuronhöz nyílt forráskódú, Trainiumra optimalizált NKI-lib is tartozik, többek között FlashAttention- és MoE-kernelekkel.

A teljesítmény elemzését a Neuron Explorer segíti. A fejlesztők ugyanazt a torch.profiler eszközt használhatják a keretrendszertől a hardverig tartó végrehajtás követésére, az eseményeket pedig a rendszer a forráskód megfelelő sorához kapcsolja.

Mit jelent ez a Trainium felhasználóinak?

A PyTorch közleménye szerint a TorchNeuron csökkenti annak szükségét, hogy a kutatók és mérnökök külön munkafolyamatot alakítsanak ki a Trainiumhoz. A meglévő eager, fordítási és elosztott PyTorch-API-k használhatók, miközben a fejlesztők egyedi NKI-kernelekkel is beavatkozhatnak a teljesítménykritikus részekbe.

A megoldás az AWS Trainium2 és Trainium3 chipeket célozza. A nagy nyelvi modellek kiszolgálására a vLLM Neuron, képek és videók generálására pedig a vLLM Omni Neuron használható beépülő modulként. A PyTorch megfogalmazása szerint a natív támogatás több hardverplatform között biztosít egységesebb fejlesztői élményt, és lehetővé teszi, hogy a csapatok a PyTorch-ökoszisztémában maradva kísérletezzenek és iteráljanak.