PyTorch 2.15: FP8 és jelentős MPS-teljesítményjavulások Macre

A Hugging Face PyTorch csapata bejelentette a PyTorch 2.15 fontos újdonságait, amelyek elsősorban az Apple MPS (Metal Performance Shaders) backendre és a Mac-en futó modellekre fókuszálnak. A kiemelt fejlesztések között szerepel az FP8 típus támogatása MPS-en, gyorsabb nyelvi modell dekódolás sűrű és Mixture-of-Experts (MoE) hálózatoknál, valamint számos MPSGraph‑ról dedikált Metal-kernelre való áttérés, ami javítja a lineáris algebrai műveletek, konvolúciók, padding és más kép‑, hang‑ és tréning műveletek teljesítményét.

Ki írta és miért fontos

A cikk szerzői a Hugging Face PyTorch csapatából származnak; aktívan hozzájárulnak a PyTorch-hoz és különösen a MPS backendre koncentrálnak, hogy a Mac-ek elsőrangú környezetté váljanak modellek futtatásához és tréningjéhez. A 2026 elején tapasztalt problémák — hiányzó operációk, teljesítmény- és memória-meglepetések — miatt a csapat valós modellekkel mért teljesítményt, és több helyen MPSGraph‑ról natív Metal kernelre váltott.

Rövid összefoglaló (TL;DR)

• FP8 mostantól elérhető MPS-en: torch.float8_e4m3fn típus, egy bájt / érték, nagyobb memóriamegtakarítás a BF16-hoz képest. Scaled matrix multiply (F.scaled_mm) támogatott.
• MoE modellek dekódolása 5–6.5× gyorsabb 2.14-hez képest: például Qwen3.6-35B-A3B M3 Ultra-n torch.compile használatával 9.4 → 61.3 tokens/s.
• Minden LLM-nél gyorsabb tokenenkénti idő: single-row linear rétegekhez specializált matrix–vector kernel, vocab‑softmax körülbelül 8× gyorsabb.
• Több lineáris algebrai művelet fut közvetlenül a GPU-n: komplex LU és triangular solve, gyorsabb kis mátrix inverziók, gyorsabb QR.
• Javított utak audiohoz, videóhoz és tréninghez: konvolúciók, padding, dropout, interpoláció backward és redukciók gyorsabbak; Wan2.2 videó VAE dekódolásnál a csúcsdriver-memória ~270 GiB-ről ~8 GiB-re csökkent.
• Több Python-vezérlés: nyilvános MPS stream-ek és megbízhatóbb esemény-időzítés.
• Új követelmények: MPS-hez macOS 15+, PyTorch 2.15-hez Python 3.11+.

Mit jelent ez a felhasználóknak?

Ha Mac-en futtatsz transformereket, diffusers‑t vagy ComfyUI‑szerű eszközöket, a gyakorlatban ezek a változások:

• Sokkal gyorsabb helyi LLM-ek: MoE modell dekódolás 5–6.5× gyorsulás (például Qwen3.6-35B-A3B és LFM2.5-8B-A1B statisztikák egy M3 Ultra-n torch.compile mellett). Sűrű modellek is profitálnak a gyorsabb lineáris rétegekből és a samplinghez szükséges softmax jelentős gyorsulásából.
• Videó-generálás, ami memóriába fér: Wan2.2 tiled VAE dekódolásnál a driver-memória csúcs ~270 GiB → ~8 GiB.
• Korrigált képszerkesztés: a Qwen-Image-2.1 VAE encoder MPS-en hibás állapotát padding-változtatás javította; most az encode–decode kör megegyezik CPU-val (PSNR 29.95 dB, korábban 10.06 dB).
• FP8 súlyok Mac-en: a torch.float8_e4m3fn típus fele annyi memóriát használ, mint BF16, és a F.scaled_mm végrehajthatja a szorzást GPU-n.

Kezdés

Szükséges hardver és szoftver: Apple silicon Mac, macOS 15 vagy újabb, Python 3.11 vagy újabb (lásd a 2.15 kompatibilitási mátrixot). Ha az OS túl régi vagy nem elérhető Metal GPU, az MPS elérhetőségi ellenőrzés megmondja az okot. A release candidate kipróbálásához külön környezetben: pip install torch==2.15.0 --index-url https://download.pytorch.org/whl/test/cpu

FP8 MPS-en

PyTorch 2.15 hozzáadja a torch.float8_e4m3fn típus támogatását MPS-re: egy bájt/érték, tehát a FP16/BF16 felének megfelelő tárhely. Ezekkel a tenzorokkal a F.scaled_mm végrehajtható, támogatva tensor‑ és rowwise skálázást is. A cikk rövid példa-kódrészletekkel mutatta be az FP8 kvantálást és a skálázott mátrixszorzást.

Gyorsabb nyelvi modell dekódolás

A dekódolási fázisban a modellek tokenenként egy sorral dolgoznak, így a nagy mátrixszorzások matrix–vector (GEMV) műveletekké szelídülnek. 2.15 több helyen optimalizálja ezt:

Linear rétegek

F.linear korábban MPSGraph-on ment; 2.15 összekapcsolja a MPS GEMV kernelt az F.linear‑lal, így a [1,1,hidden_size] alakú dekódolási input közvetlenül használhatja a gyors GEMV-t, és ahol lehet a bias hozzáadás beolvad. 96 BF16 rétegalakzat valós modellekből átlagosan ~1.23× gyorsulást mutatott; a legnagyobb nyereség a kis projekciós rétegeknél látszik (például MoE gatingek).

Mixture‑of‑Experts (MoE) rétegek

MoE rétegeknél a tokenek csoportokra oszlanak és a grouped matmul kezeli az egyenetlen csoportokat egyetlen művelettel. 2.15‑ben MPS-hez natív grouped matmul implementáció készült: az offseteket a GPU-n olvassa be, Metal Performance Primitives (MPP) használható ott, ahol elérhető, és külön Metal út a többi esethez. Kis csoportokra (gyakoribb a dekódolásnál) külön kernel érhető el. Ennek hatása a token-throughputban látszik: M3 Ultra-n torch.compile mellett MoE modellek körülbelül 1.3× több tokent/s generálnak. Konkrét példaként Qwen3.6-35B-A3B 9.4 → 61.3 tokens/s (ez a TL;DR-ben részletezett mért érték).

Lineáris algebrai műveletek a GPU-n

• Komplex számok: complex64 LU faktorizáció, LU solve és triangular solve most MPS-en futnak, ami felszabadítja a rájuk épülő műveleteket és gradiensszámítást (például komplex determinánsok).
• Kis mátrixok: FP32 mátrixok 4×4-ig egyetlen Metal kernelben invertálhatók, ami hasznos kamera-geometria és képi transzformációk esetén.
• Nagyobb feladatok: QR blokkokkal gyorsítva, triangular solve-ok munkájának nagy része hatékony mátrixszorzássá alakul.
• Gated DeltaNet: dedikált gyors út 16×16, 32×32 és 64×64 FP32 triangular rendszerekre legalább 8 jobb oldali oszlop esetén.
• SVD robosztusabb lett szélsőséges numerikus esetekben (nagyon kicsi/sok/…), így pseudoinverse és least-squares eredmények megbízhatóbbak.

Konvolúciók, padding és tréning-optimalizációk

Számos tréning- és multimédiás művelet kapott Metal-alapú gyorsabb utakat:

• conv1d forward: specializált kernellek MPSGraph helyett, plusz im2col+matmul rövid kimenetekhez; hasznos audio és streaming alakzatokhoz.
• Constant padding: Metal kernelre költözött — ez megoldotta a tiled VAE videó dekódolásnál előforduló robbanásszerű driver-memória növekedést.
• Reflection és replication padding: új Metal megvalósítások, gyorsabbak és helyesek a backward esetben.
• Out-of-place dropout: nn.Dropout és F.dropout most a meglévő fused kernelt használja.
• Upsampling backward: Metal gather kernel váltotta az MPSGraph resize-gradient útját.
• Channels-last conv2d backward: hatékony channels-last út marad, nem esik vissza NCHW‑s, költséges strided hozzáférésekre.

Softmax és redukciók

A token sampling gyorsul: softmax és log_softmax dedikált Metal kernellel futnak. Nagy szókincsek esetén a gyakorlatban egy [1,128256] FP32 softmax egy soron ~57 µs-ról ~7 µs-ra esett, kb. 8× gyorsulás. Hasonló javulások látszanak más gyakori vocab méreteknél is.

Jelentősen kisebb memóriafogyasztás videó-dekódolásnál

A constant-padding Metalre költözése különösen látványos eredményt hozott: Wan-AI/Wan2.2-TI2V-5B-Diffusers tiled VAE dekódolásnál 28-tile dekódás során 8,764 padding hívás történik; MPSGraph padding mellett a teljes driver-memória csúcs ~270.135 GiB volt, Metal padding mellett ez ~8.108 GiB‑re csökkent. A PyTorch tenzor- és allocator csúcsok viszonylag kicsik maradtak (~4.4 GiB, ~8.07 GiB).

Több vezérlés Pythonból: MPS stream-ek

Mostantól kiválaszthatod, melyik MPS stream futtatja a munkát Pythonból: torch.mps.Stream és a torch.mps.stream kontextusmenedzser használható, és Tensor.record_stream lehetővé teszi az allocator számára, hogy kövesse egy tenzor használatát más stream-en. A cikk rövid példát mutatott egy külön munkastrím használatára és szinkronizációra.

Több MPS-en futó művelet

Újonnan MPS-en futó műveletek: torch.frexp, torch.logspace, torch.take, Tensor.put_, torch.heaviside, torch.special.laguerre_polynomial_l és unsigned-integer összehasonlítások. logspace még nem támogat integer vagy boolean outputokat több elemmel. Az attention MPS-implementációján is javítottak: a scaled dot-product attention immár helyesen kezeli a batch és head broadcastolást.

Egy pillantás a 2.16-ra

A 2.16 fejlesztései már folyamatban vannak: több MPSGraph→Metal migráció, attention teljes lehúzása Metalra, és további LLM-dekódolási gyorsítások. Például a parancs-buffer commitolását 16 kernelenként tervezik, ami lehetővé teszi, hogy a GPU korábban kezdje a végrehajtást, míg a CPU még kódol. Kísérleti értékek alapján egyes modellek 4–19% gyorsulást mutattak a compiled BF16 decoding mérésénél.

Egyéb tervek: Batch norm forward/backward migrációk, MSE loss, std_mean/var_mean és bincount továbbfejlesztései (bincount egy példa: 4 millió integer megszámlálása korábban 6.8 s volt, most ~0.5 ms, ha a legtöbb érték egy binbe esik).

Záró gondolatok

A PyTorch 2.15 jelentős előrelépés a Mac-en történő modellezés és tréning szempontjából: FP8 támogatás, gyorsabb dekódolás MoE és sűrű LLM-eknél, több lineáris algebrai művelet GPU-n és sok MPSGraph → Metal migráció, amelyek összességében jobb teljesítményt és alacsonyabb memóriahasználatot eredményeznek. A csapat tovább dolgozik a 2.16‑on és szívesen várnak visszajelzéseket a GitHubon.

Köszönetnyilvánítás: a cikk megemlíti Nikita Shulga, Kurt Mohler és Aaron Gokaslan munkáját és a számos közreműködőt, akik a MPS fejlesztéséhez hozzájárultak a kiadásban.