Az AMD GPU-kon is gyorsult az FP8-as PyTorch-tréning
A PyTorch közvetlen AMD Instinct GPU-támogatással bővítette a TorchTitan rendszert, miközben a TorchAO FP8-optimalizációi több modellen is gyorsulást hoztak. Az AMD fejlesztései bekerültek a nyilvános pytorch/AO és pytorch/TorchTitan projektekbe.
- A TorchTitan közvetlen AMD Instinct GPU-támogatást kapott.
- A Llama3-8B FP8-tréningje 13,4 százalékkal gyorsabb volt a BF16-nál.
- A DeepSeek-MoE-16B visszafelé futó számítása 4,2-szer gyorsult.
- A DeepSeek-V3 671B egyes MoE-rétegei 6,2-szer gyorsultak.
- A Triton-optimalizáció a DeepSeek-V3 671B FP8-különbségének 89 százalékát ledolgozta.
Az AMD FP8-formátumának támogatása
A PyTorch 2026. augusztus 13-i bejegyzése szerint az AMD Instinct GPU-khoz készült optimalizációkat a TorchTitan és a TorchAO upstream kódjába is beolvasztották. A TorchTitan így közvetlenül támogatja az AMD Instinct GPU-kat, a PyTorch szerint versenyképes FP8-teljesítménnyel.
Az FP8-tréning a lineáris rétegek három mátrixszorzását, a forward lépést, a gradiens bemenetének számítását és a súlygradiens frissítését 16 bites formátumról 8 bitesre alakítja. Az AMD Instinct GPU-k az FP8 e4m3fnuz változatát használják. Ennél a maximális érték 240, és nincs külön NaN vagy végtelen érték kódolás.
A formátum helyes felismerése a számítási pontosság feltétele. A TorchAO korábban más maximális értékkel számolta a skálázást, ezért az AMD hardverén az aktivációk és gradiensek értéktartománya hibás lehetett. A fejlesztők automatikus hardverfelismerést vezettek be, így a könyvtár a megfelelő FP8-típust és maximális értéket választja ki.
Gyorsulás sűrű és MoE-modelleken
A sűrű Llama3-8B modellen, 8 darab MI300X GPU használatával, a soronkénti FP8-tréning 13,4 százalékkal nagyobb áteresztőképességet ért el a BF16-hoz képest. A mérésben a súlyfrissítéshez továbbra is BF16-ot használtak, miközben a forward és a gradiensbemeneti mátrixszorzások FP8-ban futottak. A csúcsmemória közel azonos maradt, körülbelül 39 GB volt.
A kevert szakértős, vagyis MoE-modelleknél összetettebb a helyzet, mert az egyes tokenek csak bizonyos szakértőkhöz kerülnek, ezért változó méretű kötegeket kell feldolgozni. A fejlesztők ROCm alatt engedélyezték az FP8-as csoportosított GEMM-et, valamint a Composable Kernel backendet használták a megfelelő adattípus és feldolgozási útvonal kiválasztására.
DeepSeek-MoE-16B esetén a visszafelé futó számítások összeolvasztása 4,2-szeres gyorsulást hozott 8 darab MI300X GPU-n. DeepSeek-V3 671B-nél egy oszloponkénti skálázási optimalizáció az egyes MoE-rétegek feldolgozási idejét 7290 mikroszekundumról 1170 mikroszekundumra csökkentette, ami 6,2-szeres gyorsulás.
A Triton-kernelok csökkentették a kvantálás költségét
Az FP8-kvantálás több különálló lépésből áll: kiszámítja a sor- vagy oszloponkénti abszolút maximumot, meghatározza a skálázási tényezőt, majd korlátozza és FP8-ra alakítja az adatot. Korábban minden lépés külön kernelindítást és köztes adatok HBM-be írását igényelte. A fejlesztők ezeket a műveleteket Triton-kernelokkal vonták össze, így kevesebb adatmozgatásra és kernelindításra volt szükség.
DeepSeek-V3 671B esetén, 8 darab MI325X GPU-n, a forward útvonal összeolvasztása 5996-ról 7027 token/másodpercre növelte a teljes áteresztőképességet, ami 17 százalékos javulás. A PyTorch szerint a módosítás az FP8 és BF16 közötti teljesítménykülönbség 89 százalékát visszahozta.
A TorchAO és a TorchTitan mind a négy bemutatott FP8-skálázási stratégiát támogatja: a tenzoronkénti, a soronkénti, a blokkonkénti és az adatok mellé csoportosan tárolt skálázást. Az AMD MI300 és MI350 GPU-khoz blokkonkénti kernel támogatása is bekerült.
PyTorch: FP8 Training on AMD GPUs with TorchTitan and TorchAO: Upstreaming Performance Improvements
