Resumo Técnico: Meganeura – Treinamento e Inferência de GPU Portáteis através de Vulkan e Metal

Declaração do Problema

O ecossistema atual de aprendizado de máquina sofre de uma fragmentação entre treinamento e implantação. O treinamento normalmente depende de frameworks Python dinâmicos (ex: PyTorch) e bibliotecas específicas de fornecedores (CUDA, ROCm), enquanto a implantação muitas vezes exige a exportação de modelos para runtimes separados (TensorRT, Core ML, motores móveis). Essa "pilha dividida" (split stack) exige trabalho duplicado em conversão de modelos, cobertura de operadores, validação de política numérica e depuração. Além disso, os modelos implantados são geralmente congelados; adaptar modelos no dispositivo exige uma ida e volta através de uma pilha de treinamento separada, impedindo o aprendizado contínuo e in loco.

O artigo investiga se um compilador nativo compacto pode abranger tanto as fases de treinamento quanto de inferência em diversos GPUs de consumo usando APIs de computação gráfica de propósito geral (Vulkan e Metal) em vez de pilhas de computação proprietárias.

Metodologia

O autor apresenta o Meganeura, um sistema nativo baseado em Rust que unifica inferência e treinamento sob um único grafo estático tipado.

• Arquitetura: O sistema utiliza um grafo tipado de nível intermediário onde os nós representam operações como multiplicação de matrizes, convolução e atenção. Ele emprega diferenciação automática de modo reverso para gerar grafos de retropropagação (backward graphs).
• Pipeline de Compilação: O grafo é reescrito (usando reescrita gananciosa e saturação de igualdade opcional), especializado para capacidades específicas de hardware (ex: matrizes cooperativas, geometria de tiles) e reduzido para WGSL (WebGPU Shading Language).
• Backend: O WGSL é validado e traduzido pelo Naga e executado via Blade, uma camada de abstração gráfica nativa que visa Vulkan (Linux, Windows, Android) e Metal (Apple).
• Execução: O sistema gera uma sequência de despacho estático com propriedade explícita de buffers e planejamento de memória. Ele não depende de um runtime Python ou captura de grafo dinâmica em tempo de execução.
• Protocolo de Avaliação: O autor utilizou o Inferena, um harness cross-framework, para comparar o Meganeura contra o PyTorch em cinco dispositivos (NVIDIA RTX 5070, AMD RX 7900 XT, AMD Radeon 780M APU, Intel iGPU, Apple M3) e cinco cargas de trabalho (SmolLM2, SmolVLA, Diffusion U-Net, ResNet-50, Whisper-tiny).
• Portões de Correção: Os resultados foram validados contra contratos numéricos estritos:
  • Forward: Erro L2 relativo < 1%.
  • Backward: Erro da norma do gradiente total < 5% e erro da norma do gradiente por parâmetro < 5%.
  • Dois modos aritméticos foram testados: Strict f32 (precisão controlada) e Accelerated (permitindo TF32 na NVIDIA e matrizes cooperativas f16 onde suportado).

Principais Contribuições

1. Pilha Unificada de Treinamento-para-Implantação: Um sistema executável compacto (34,5 KLOC de Rust, 6,1 KLOC de WGSL) que suporta inferência, diferenciação de modo reverso, atualizações SGD/Adam e transferência de checkpoint sem dependências de runtime Python ou CUDA/ROCM.
2. Avaliação Comparativa Controlada: Uma comparação rigorosa através de 50 células de dispositivo-carga de trabalho-modo. 48 células passaram tanto no portão forward quanto no backward. O estudo isola as lacunas de desempenho em famílias de kernels específicos em vez de atribuí-las às APIs gráficas em si.
3. Decomposição de Desempenho: O artigo atribui as diferenças de desempenho a fatores concretos: cobertura de kernel (especificamente derivadas de convolução e backward de atenção), estrutura de despacho e política aritmética (seleção de matriz cooperativa), em vez de limitações inerentes da API.
4. Estudo de Caso de Implantação Física: Uma demonstração do DinoVision, uma aplicação Android XR que compartilha uma fila gráfica entre a renderização e um decoder treinado pelo Meganeura, validando a capacidade da pilha de coexistir com cargas de trabalho gráficas.

Resultados

Portabilidade de Desempenho

• Dispositivos AMD: O Meganeura alcançou paridade ou melhor desempenho comparado ao PyTorch compilado com ROCm. No RX 7900 XT, quatro das cinco cargas de trabalho de inferência ficaram dentro de 1,10× do PyTorch, e três cargas de treinamento foram mais rápidas.
• Dispositivos NVIDIA: No modo strict f32, o Meganeura venceu 12 de 20 células de latência mínima. No entanto, as lacunas de treinamento foram maiores (até 4,6× no ResNet-50), principalmente devido à ausência de kernels de derivada de convolução otimizados e à falta de suporte a TF32 no caminho do Meganeura.
• Apple Silicon: A latência de inferência foi competitiva (0,93–1,88×), mas o treinamento ficou atrás (1,20–2,84×) em relação ao MPS eager, largely devido a ineficiências no kernel de backward de atenção.
• Intel iGPU: Como o caminho XPU do fornecedor não estava funcional no ambiente de teste, o Meganeura superou o fallback de CPU em 12 de 15 células.
• Velocidade de Compilação: O Meganeura compilou cargas de trabalho em 0,1–2,4 segundos, significativamente mais rápido que o torch.compile (6–96 segundos).

Correção e Estabilidade

• Taxa de Sucesso: 48 de 50 células passaram pelos portões forward e backward.
• Anomalias: Duas falhas ocorreram no backward pass do Whisper no AMD 780M APU. A análise de gradiente entre dispositivos sugeriu que o caminho de referência (PyTorch/ROCM) era a fonte do erro, não o Meganeura, embora uma terceira implementação seja necessária para resolução.
• Precisão Numérica: No modo strict f32, o pior erro L2 forward foi 0,0043% (portão: 1%), e o pior erro de norma de gradiente foi 0,64% (portão: 5%).

Sistemas e Implantação

• Tamanho do Artefato: O binário stripped do Meganeura tem 13 MiB, comparado a ~4,66 GiB para um fechamento completo do runtime PyTorch CUDA.
• Memória: O planejador de memória estática do Meganeura reduziu a alocação física em 4–48% no treinamento e 6–79% na inferência em comparação aos totais lógicos, embora alocadores dinâmicos no PyTorch tenham alcançado picos menores para cargas de trabalho específicas com muitas ativações.
• Co-tenência: O estudo de caso DinoVision demonstrou que o Meganeura pode compartilhar uma fila Vulkan com a renderização OpenXR, reduzindo a latência do worker em 20,2% ao usar dois chunks em comparação com um.

Significância e Alegações

O artigo afirma que APIs gráficas de consumo geral (Vulkan e Metal) são substratos suficientes para uma pilha compacta e compartilhada de treinamento-para-implantação que alcança desempenho útil e, às vezes, competitivo com os fornecedores.

• Não é um Substituto para Pesquisa: O autor afirma explicitamente que o Meganeura não é um substituto para fluxos de pesquisa Python dinâmicos, treinamento em escala de nuvem distribuída ou a amplitude total de operadores de frameworks maduros como o PyTorch.
• Lacunas Específicas Identificadas: As lacunas de desempenho restantes são atribuídas à cobertura de kernels (especificamente derivadas de convolução na NVIDIA e backward de atenção no Metal) e políticas de escalonamento (seleção de matriz cooperativa), não às APIs gráficas em si.
• Portabilidade vs. Desempenho: Os resultados sugerem que, embora uma única implementação possa não igualar o desempenho de pico do fornecedor em cada configuração de hardware, ela oferece portabilidade de desempenho superior para latência de lote mínimo e reduz significativamente o overhead de engenharia de manter pilhas de treinamento e implantação separadas.
• Adaptação ao Vivo: O trabalho estabelece o substrato técnico necessário para a "adaptação ao vivo" em dispositivos implantados, onde modelos podem ser ajustados ou personalizados in-place sem sair do código-fonte da aplicação.

O artigo conclui que o problema da "pilha dividida" pode ser resolvido por um compilador nativo unificado, deslocando o fardo da engenharia da gestão de múltiplas fronteiras de conversão para a otimização de um único compilador e artefato inspecionável.