Technische Samenvatting: Meganeura – Draagbare GPU-training en -inferentie via Vulkan en Metal

Probleemstelling

Het huidige machine learning-ecosysteem lijdt aan een fragmentatie tussen training en deployment. Training leunt doorgaans op dynamische Python-frameworks (bijv. PyTorch) en leveranciersspecifieke bibliotheken (CUDA, ROCm), terwijl deployment vaak vereist dat modellen naar aparte runtimes wordt geëxporteerd (TensorRT, Core ML, mobiele engines). Deze "gesplitste stack" noodzaakt dubbel werk bij modelconversie, dekking van operatoren, validatie van numerieke policies en debugging. Bovendien zijn uitgerolde modellen meestal statisch; het aanpassen van deze op het apparaat zelf vereist een retourgang via een aparte trainingstack, wat continue, in-place leerprocessen verhindert.

Dit onderzoek onderzoekt of een enkele, compacte native compiler zowel de trainings- als de inferentiefase kan overbruggen over diverse consumenten-GPU's met behulp van algemene graphics-API's (Vulkan en Metal) in plaats van propriëtaire compute-stacks.

Methodologie

De auteur presenteert Meganeura, een native Rust-gebaseerd systeem dat inferentie en training verenigt onder één enkele getypeerde statische graaf.

• Architectuur: Het systeem gebruikt een mid-level getypeerde graaf waarbij nodes operaties vertegenwoordigen zoals matrixvermenigvuldiging, convolutie en attention. Het maakt gebruik van reverse-mode automatic differentiation om backward-grafen te genereren.
• Compilatiepijplijn: De graaf wordt herschreven (via greedy rewriting en optionele equality saturation), gespecialiseerd voor specifieke hardwarecapaciteiten (bijv. cooperative matrices, tile geometry), en verlaagd naar WGSL (WebGPU Shading Language).
• Backend: WGSL wordt gevalideerd en vertaald door Naga en uitgevoerd via Blade, een native graphics-abstractielaag die zich richt op Vulkan (Linux, Windows, Android) en Metal (Apple).
• Executie: Het systeem genereert een statische dispatch-sequentie met expliciete buffer-eigendom en geheugenplanning. Het vertrouwt niet op een Python-runtime of dynamische graaf-capture tijdens runtime.
• Evaluatieprotocol: De auteur gebruikte Inferena, een cross-framework harness, om Meganeura te vergelijken met PyTorch over vijf apparaten (NVIDIA RTX 5070, AMD RX 7900 XT, AMD Radeon 780M APU, Intel iGPU, Apple M3) en vijf workloads (SmolLM2, SmolVLA, Diffusion U-Net, ResNet-50, Whisper-tiny).
• Correctheidspoorten: De resultaten werden gevalideerd tegen strikte numerieke contracten:
  • Forward: Relatieve L2-fout < 1%.
  • Backward: Totale gradiënt-norm fout < 5% en per-parameter gradiënt-norm fout < 5%.
  • Twee rekenkundige modi werden getest: Strict f32 (gecontroleerde precisie) en Accelerated (toestaan van TF32 op NVIDIA en f16 cooperative matrices waar ondersteund).

Belangrijkste Bijdragen

1. Unified Train-to-Deploy Stack: Een compact executable systeem (34,5 KLOC Rust, 6,1 KLOC WGSL) dat inferentie, reverse-mode differentiatie, SGD/Adam updates en checkpoint-transfer ondersteunt zonder afhankelijkheid van een Python-runtime of CUDA/ROCm-dependencies.
2. Gecontroleerde Vergelijkende Evaluatie: Een rigoureuze vergelijking over 50 device-workload-mode cellen. 48 cellen slaagden voor zowel de forward als de backward correctheidspoorten. De studie isoleert prestatiekloven tot specifieke kernelfamilies in plaats van ze toe te schrijven aan de graphics-API's zelf.
3. Prestatiedecompositie: De paper schrijft prestatieverschillen toe aan concrete factoren: kerneldekking (specifiek convolutie-derivaten en attention backward), dispatch-structuur en de rekenkundige policy (cooperative matrix selectie), in plaats van aan inherente API-beperkingen.
4. Fysieke Deployment Casestudy: Een demonstratie van DinoVision, een Android XR-applicatie die een graphics queue deelt tussen rendering en een door Meganeura getrainde decoder, wat de capaciteit van de stack valideert om naast graphics-workloads te co-existeren.

Resultaten

Prestatieportabiliteit

• AMD Devices: Meganeura bereikte een prestatie die bijna gelijkwaardig is aan of beter is dan gecompileerde ROCm PyTorch. Op de RX 7900 XT lag vier van de vijf inferentie-workloads binnen 1,10× van PyTorch, en drie training-workloads waren sneller.
• NVIDIA Devices: In strict f32-modus won Meganeura 12 van de 20 minimal-latency cellen. Echter, de gaten bij training waren groter (tot 4,6× op ResNet-50), primair door ontbrekende geoptimaliseerde convolutie-derivaat kernels en het gebrek aan TF32-ondersteuning in het Meganeura-pad.
• Apple Silicon: Inferentie-latentie was competitief (0,93–1,88×), maar training liep achter (1,20–2,84×) ten opzichte van eager MPS, grotendeels door inefficiënties in de backward attention kernels.
• Intel iGPU: Omdat het vendor XPU-pad niet functioneel was in de testomgeving, presteerde Meganeura beter dan de CPU-fallback in 12 van de 15 cellen.
• Compilatiesnelheid: Meganeura compileerde workloads in 0,1–2,4 seconden, aanzienlijk sneller dan torch.compile (6–96 seconden).

Correctheid en Stabiliteit

• Succespercentage: 48 van de 50 cellen slaagden voor zowel de forward als de backward poorten.
• Anomalieën: Twee fouten traden op bij de AMD 780M APU tijdens de Whisper backward pass. Cross-device gradiëntanalyse suggereerde dat het referentiepad (PyTorch/ROCm) de bron van de fout was, en niet Meganeura, hoewel een derde implementatie nodig is voor resolutie.
• Numerieke Precisie: In strict f32-modus was de slechtste forward L2-fout 0,0043% (poort: 1%), en de slechtste gradiënt-norm fout was 0,64% (poort: 5%).

Systemen en Deployment

• Artifact Grootte: De gestripte Meganeura-binary is 13 MiB, vergeleken met ~4,66 GiB voor een volledige PyTorch CUDA runtime closure.
• Geheugen: De statische geheugenplanner van Meganeura verminderde de fysieke allocatie met 4–48% in training en 6–79% in inferentie vergeleken met de logische totalen, hoewel dynamische allocators in PyTorch soms lagere pieken bereikten voor specifieke activatie-zware workloads.
• Co-tenancy: De DinoVision casestudy toonde aan dat Meganeura een Vulkan queue kan delen met OpenXR rendering, wat de worker-latentie met 20,2% vermindert wanneer twee chunks worden gebruikt in vergelijking met één.

Betekenis en Claims

De paper claimt dat algemene consumenten-graphics API's (Vulkan en Metal) voldoende substraten zijn voor een compacte, gedeelde train-to-deploy stack die een bruikbare, en soms vendor-competitieve, prestatie levert.

• Geen Vervanging voor Research: De auteur stelt expliciet dat Meganeura geen vervanging is voor dynamische Python research workflows, gedistribueerde cloud-scale training, of de volledige operator-breedte van volwassen frameworks zoals PyTorch.
• Specifieke Gaten Geïdentificeerd: De resterende prestatiekloven worden toegeschreven aan kerneldekking (specifiek convolutie-derivaten op NVIDIA en attention backward op Metal) en scheduling policies (cooperative matrix selectie), en niet aan de graphics-API's zelf.
• Portabiliteit vs. Prestatie: De resultaten suggereren dat hoewel een enkele implementatie de piek-vendorprestaties op elke hardwareconfiguratie mogelijk niet evenaart, het een superieure prestatieportabiliteit biedt voor minimal-batch latentie en de engineering-overhead voor het onderhouden van aparte training- en deploymentstacks aanzienlijk vermindert.
• Live Adaptatie: Het werk legt het technische substraat vast dat vereist is voor "live adaptatie" op uitgerolde apparaten, waarbij modellen in-place kunnen worden gefinetuned of gepersonaliseerd zonder de applicatiecode te verlaten.

De paper concludeert dat het "split stack"-probleem kan worden opgelost door een unificerende native compiler, waarbij de engineering-last verschuift van het beheren van meerdere conversiegrenzen naar het optimaliseren van een enkele, inspecteerbare compiler en artifact.