Gyorsabb LLM-ek és FP8 támogatás érkezik a PyTorch 2.15 MPS backendjébe

FP8 és gyorsabb nyelvimodell-dekódolás

A Hugging Face PyTorch-csapata október 20-án ismertette a PyTorch 2.15 MPS-fejlesztéseit. Az MPS a PyTorch Apple-féle Metal-alapú gyorsítási backendje, amely Mac gépeken használható modellek futtatására és tanítására.

Az új kiadásban az MPS támogatja a torch.float8_e4m3fn tenzorokat. Az FP8-értékek bájtonként egy bájtot foglalnak, így a BF16-hoz és az FP16-hoz képest fele akkora tárhelyet igényelnek. Az értékek az F.scaled_mm művelettel közvetlenül a GPU-n szorozhatók, tenzoronkénti és soronkénti skálázással.

A nyelvi modellek dekódolási szakaszában a PyTorch 2.15 külön optimalizálja a sűrű lineáris rétegeket és a mixture-of-experts, vagyis MoE-rétegeket. A lineáris rétegek 96, valódi modellekből származó BF16-alakzatán mérve a geometriai átlagos gyorsulás 1,23-szoros. Az MoE-modellekhez natív csoportos mátrixszorzás érkezik, amely csökkenti a GPU és a CPU közötti szinkronizációt.

Nagy ugrás az MoE-modellek és a videó memóriahasználatában

A Hugging Face mérései szerint az MoE-modellek dekódolása a PyTorch 2.14-hez képest 5 és 6,5-szer gyorsabb lehet. M3 Ultra gépen, torch.compile használatával a Qwen3.6-35B-A3B sebessége 9,4-ről 61,3 token/másodpercre nőtt, az LFM2.5-8B-A1B pedig 23,6-ról 123,6 token/másodpercre gyorsult.

A sűrű modellek is profitálnak a gyorsabb lineáris rétegekből. A szókincsen végzett softmax művelet körülbelül nyolcszor gyorsabb, ami a mintavétellel dolgozó dekódolási folyamatot is érinti.

A kiadás a videógenerálási munkafolyamatok memóriaigényét is mérsékli. A Wan2.2 csempézett VAE-dekódolása korábban 270 GiB-re növelhette a meghajtó memóriáját, az új megoldásban a csúcsérték körülbelül 8 GiB. A Qwen-Image-2.1 VAE-kódolójának MPS-en jelentkező képhibája szintén javult: egy kódolásból és dekódolásból álló körben a CPU-eredményhez viszonyított PSNR 10,06 dB-ről 29,95 dB-re emelkedett.

Több művelet fut Metal-kernellel

A PyTorch 2.15 további lineáris algebrai műveleteket támogat MPS-en. A complex64 típusú LU-felbontás, LU-megoldás és háromszögmátrixos megoldás már a GPU-n is futtatható. Gyorsabb útvonal érkezik a legfeljebb 4×4-es FP32-mátrixok kötegelt invertálásához, a QR-felbontáshoz és bizonyos háromszögmátrixos feladatokhoz is.

Az audio-, videó- és tanítási munkafolyamatokban több művelet kapott dedikált Metal-megoldást. A fejlesztők külön kiemelik a konvolúciókat, a kitöltést, a dropoutot, az interpoláció visszafelé irányuló műveletét és a redukciókat. A Hugging Face szerint a fejlesztések jelentős része az MPSGraph helyett közvetlen Metal-kernellel valósul meg.

A használathoz Apple chippel szerelt Mac, macOS 15 vagy újabb rendszer és Python 3.11 vagy újabb szükséges. A projekt már a PyTorch 2.16 irányába is előretekint: a tervek között szerepel az attention teljes Metalra helyezése, további MPSGraph és Metal közötti átállás, valamint a nyelvi modellek dekódolásának újabb gyorsítása.

Mit jelent ez a Mac-felhasználóknak?

A frissítés azoknak lehet különösen fontos, akik Mac gépen futtatnak modelleket a transformers, a diffusers vagy például a ComfyUI segítségével. Az FP8 kisebb memóriaterhelést tesz lehetővé, az MoE-modellek gyorsabb dekódolása rövidebb tokenenkénti várakozást eredményezhet, a Wan2.2 esetében pedig a korábbinál jóval alacsonyabb memóriaigény mellett válhat lehetővé a videódekódolás.

A PyTorch 2.15 emellett több korábban problémás művelethez kínál MPS-es útvonalat, így a Hugging Face célja szerint a Mac alkalmasabbá válik PyTorch-modellek futtatására és tanítására.