Trénink LLM ve Swiftu: od 2,8 Gflop/s na 11 Tflop/s — jak na to?
Článek na Cocoa with Love ukazuje 382násobné zrychlení maticového násobení ve Swiftu na Apple Silicon — od naivní implementace přes SIMD až po Metal GPU shadery.
Blog Cocoa with Love publikoval detailní průvodce optimalizací maticového násobení ve Swiftu pro trénink jazykových modelů na Apple Silicon. Výsledek: 382násobné zrychlení a finální výkon 11,1 Tflop/s na GPU.
Výchozí bod a cíl
Naivní implementace maticového násobení ve Swiftu dosahuje 2,8 Gflop/s. Cíl článku je dostat se přes hranici 1 Tflop/s — a pak dál. Postupuje se po vrstvách optimalizací.
Sedm technik optimalizace
1. MutableSpan
Eliminace Copy-on-Write režie u polí. Swift pole jsou value types — při každé mutaci hrozí kopírování. MutableSpan to obchází.
2. Relaxed.multiplyAdd
Využití FMA instrukcí (fused-multiply-add) z knihovny Swift-Numerics. Jedna instrukce místo dvou pro a * b + c.
3. InlineArray
Náhrada dynamických polí statickými poli na stacku. Méně heap alokací = lepší cache lokalita.
4. Loop unrolling
Rozvinutí vnitřních smyček pro lepší SIMD zpracování. Kompilátor pak snáze vektorizuje výpočty.
5. Multithreading
Paralelizace řádků matice pomocí DispatchQueue.concurrentPerform. Lineární škálování s počtem jader.
6. AMX jednotka
Přímý přístup k Apple Matrix Coprocessoru (AMX) — dedikovaný hardware pro maticové operace přítomný v každém Apple Silicon čipu.
7. Metal GPU shadery
Compute shadery s tile-based optimalizací využívající sdílenou paměť GPU. Tady se dosahuje finálních 11,1 Tflop/s.
Výsledky podle metody
| Metoda | Výkon |
|---|---|
| Naivní Swift | 2,8 Gflop/s |
| SIMD + FMA | ~50 Gflop/s |
| AMX | ~1,1 Tflop/s |
| Metal GPU | 11,1 Tflop/s |
Zajímavý postřeh: Swift ve finální implementaci dosahuje srovnatelného nebo vyššího výkonu než referenční implementace v C.