Orion: Characterizing and Programming Apple's Neural Engine for LLM Training and Inference

O artigo apresenta o Orion, o primeiro sistema de ponta a ponta que permite o treinamento e inferência estáveis de modelos de linguagem diretamente na Neural Engine da Apple, contornando as limitações do CoreML através de APIs privadas, descobrindo novas restrições de hardware e otimizando o processo de treinamento ao reduzir drasticamente o tempo de recompilação de pesos.

O essencial

Imagine que você tem um carro de luxo (seu iPhone ou Mac) que vem de fábrica com um motor de corrida secreto escondido no porta-malas. Esse motor é superpotente, feito especificamente para correr em pistas de inteligência artificial. Ele é chamado de Neural Engine (Motor Neural) da Apple.

O problema é que, até agora, ninguém sabia como ligar esse motor. A Apple dizia: "Use o motor principal (CPU) ou o motor gráfico (GPU) para rodar seus aplicativos". O motor de corrida ficava lá, frio e parado, enquanto o carro usava um motor a diesel para fazer o trabalho pesado de entender linguagem humana (como o ChatGPT).

O artigo "Orion" é como um manual de instruções pirata e genial que ensina como:

1. Desbloquear esse motor secreto.
2. Consertar os problemas que faziam o carro engasgar.
3. Fazer o carro acelerar em treinamento de inteligência artificial sem precisar trocar o motor a cada curva.

Aqui está a explicação do que eles fizeram, usando analogias do dia a dia:

1. O Problema: A "Caixa Preta"

A Apple tem um sistema chamado CoreML que é como um gerente de trânsito muito rígido. Ele decide sozinho para onde o tráfego (os dados) deve ir. Os programadores não podiam mandar o carro usar o motor de corrida (Neural Engine) diretamente. Era como tentar dirigir um carro de F1, mas o volante estava trancado e o gerente de trânsito dizia: "Vou decidir se você usa as rodas da frente ou de trás".

Além disso, treinar uma IA (ensinar o carro a dirigir sozinho) exigia reescrever o manual de instruções do motor a cada vez que o carro aprendia algo novo. Isso era tão lento que o carro ficava parado na garagem a maior parte do tempo.

2. A Solução: O Projeto Orion

Os autores criaram o Orion, um sistema que ignora o gerente de trânsito e vai direto ao motor. Eles usaram "chaves mestras" (APIs privadas) para falar diretamente com o hardware.

A Grande Descoberta: "Troca de Peso Cirúrgica" (Delta Compilation)

Esta é a parte mais brilhante do artigo.

• O jeito antigo (Lento): Imagine que você está treinando um atleta. A cada treino, você precisava derreter a estátua de bronze do atleta, esculpir uma nova com os músculos atualizados e fundi-la novamente antes de ele poder correr. Isso levava 4 segundos por treino. Se você fizesse 1.000 treinos, passaria 4.000 segundos apenas fundindo estátuas!
• O jeito Orion (Rápido): Eles descobriram que não precisavam derreter a estátua inteira. Eles podiam apenas trocar os pesos (as roupas e acessórios) do atleta enquanto ele estava pronto.
  • Eles "desligam" o programa antigo do motor.
  • Trocam os arquivos de peso no disco (como trocar a camisa do atleta).
  • "Ligam" o programa de novo.
  • Resultado: O que levava 4 segundos agora leva 0,5 segundos. O treino ficou 3,8 vezes mais rápido.

O "Motor de Corrida" e as Regras Escondidas

O motor de corrida (Neural Engine) é muito rápido, mas tem manias estranhas, como um carro de F1 que só aceita combustível de uma marca específica e só pode fazer curvas em um ângulo exato.
Os pesquisadores descobriram 20 regras secretas que a Apple não documentou. Por exemplo:

• Se você tentar juntar duas peças de dados de um jeito errado, o motor para de funcionar.
• Se a memória for muito grande, o motor "engasga".
• Eles tiveram que criar um tradutor (Compilador) que pega as instruções gerais de IA e as traduz para a linguagem secreta que o motor entende, garantindo que nada quebre.

3. O Que Conseguiram Fazer?

Com o Orion, eles provaram que é possível:

• Treinar IA no seu Mac: Eles treinaram um modelo de linguagem (como um mini-ChatGPT) diretamente no chip do computador, sem usar a nuvem.
• Estabilidade: Antes, o motor "engasgava" e a IA ficava louca (erros numéricos). Eles consertaram três bugs que causavam isso, garantindo que o treinamento fosse estável por 1.000 passos sem travar.
• Troca Rápida de "Adaptadores" (LoRA): Imagine que você quer que o motorista saiba dirigir um caminhão, depois um barco, depois um avião. Antigamente, você teria que trocar o motor inteiro. Com o Orion, você apenas troca o "volante" (os adaptadores) instantaneamente, sem parar o motor.

4. Por que isso é importante?

Hoje, temos 2 bilhões de dispositivos Apple (iPhones, iPads, Macs) com esse motor de corrida superpotente, mas ele está dormindo.

• Economia de Energia: O motor de corrida gasta quase zero energia quando está parado e é super eficiente.
• Privacidade: Você pode treinar sua própria IA no seu celular, sem enviar seus dados para a nuvem da Apple ou de outras empresas.
• Acesso: Como o Orion é de código aberto (grátis), qualquer desenvolvedor pode usar esse motor secreto para criar novas tecnologias.

Resumo em uma frase

O Orion é como um mecânico genial que descobriu como ligar o motor de F1 secreto da Apple, consertou os vazamentos de combustível e criou um sistema para trocar as peças do carro enquanto ele ainda está em movimento, permitindo que qualquer pessoa treine inteligências artificiais super-rápidas e privadas diretamente no seu próprio dispositivo.

Resumo técnico

Visão Geral

O artigo apresenta o Orion, o primeiro sistema aberto de ponta a ponta capaz de programar diretamente o Apple Neural Engine (ANE) para treinamento e inferência de Grandes Modelos de Linguagem (LLMs). Enquanto o ANE está presente em mais de dois bilhões de dispositivos Apple, ele permanece subutilizado para LLMs devido às abstrações opacas do CoreML, que impedem o controle direto e o treinamento. O Orion contorna o CoreML, utilizando APIs privadas para criar um runtime nativo, um compilador e uma técnica inovadora de "compilação delta" que viabiliza o treinamento estável no dispositivo.

---

1. O Problema

• Subutilização de Hardware: O ANE (presente desde o A11 Bionic até o M4) oferece alto desempenho (até 38 TOPS em INT8, ~19 TFLOPS em FP16 medidos), mas frameworks públicos como CoreML atuam como "caixas pretas", decidindo dinamicamente onde executar operações (CPU, GPU ou ANE) sem permitir que o desenvolvedor force a execução no ANE.
• Falta de Suporte a Treinamento: Não existe nenhum framework público que suporte treinamento de LLMs no ANE. O CoreML não suporta computação de gradientes.
• Barreiras de Programação: O conjunto de instruções nativo do ANE (MIL - Model Intermediate Language) é não documentado. As APIs de compilação e avaliação residem em um framework privado (AppleNeuralEngine.framework).
• Limitações de Estado Anterior: Trabalhos anteriores (como maderix e ANEgpt) desvendaram as APIs privadas e caracterizaram o hardware, mas falharam em criar um sistema completo: ANEgpt sofria de divergência numérica (NaN) ao retomar o treinamento, não possuía compilador e dependia de Python para orquestração.

---

2. Metodologia e Arquitetura

O Orion é estruturado em cinco camadas, operando inteiramente em Objective-C (com Python apenas para conversão inicial de pesos), bypassando o CoreML:

1. Caracterização do ANE: Os autores expandiram o conhecimento sobre o hardware, catalogando 20 restrições de programação (6 conhecidas, 14 novas descobertas). Estas incluem:

   • Restrições de IR MIL: Operações como concat são rejeitadas; gelu deve ser substituído por uma aproximação tanh.
   • Layout de Memória: Requisitos estritos de alinhamento e ordenação alfabética de buffers de entrada/saída.
   • Limites de Compilação: Um limite de ~119 compilações por processo antes de falhas silenciosas.
   • Comportamento Numérico: O ANE "assada" (bakes) os pesos no momento da compilação, tornando-os imutáveis após o carregamento.

2. Pipeline de Compilação:

   • Um compilador que transforma um IR de grafo (27 operações) em código MIL nativo do ANE.
   • Passes de otimização: Eliminação de código morto, fusão de casts, anotação de SRAM e validação de restrições.
   • Suporte a frontends verificados para GPT-2 e modelos de treinamento (Stories110M).

3. Runtime e I/O:

   • Gerenciamento de memória zero-copy usando IOSurface para transferência de tensores entre CPU e ANE.
   • Cache de programas compilados para evitar recompilações desnecessárias.

4. Divisão de Trabalho (CPU/ANE):

   • ANE: Executa operações computacionalmente intensivas e limitadas (convoluções 1x1, camadas de atenção, FFN).
   • CPU: Executa amostragem de tokens, otimizador Adam, acúmulo de gradientes e operações não suportadas (como gather ou camadas com canais muito altos).

---

3. Contribuições Chave

A. Compilação Delta (Delta Compilation)

Esta é a contribuição mais significativa para o treinamento.

• O Desafio: Como o ANE "assada" os pesos na compilação, qualquer atualização de peso exigiria recompilar todo o programa. No modelo anterior (v1.0), isso levava ~4.2 segundos por passo de treinamento (83.9% do tempo total) e exigia reiniciar o processo para contornar o limite de compilações.
• A Solução: O Orion utiliza uma técnica de "recarregamento cirúrgico". Em vez de recompilar, o sistema:
  1. Descarrega o programa do ANE (unload).
  2. Atualiza os arquivos de pesos no disco (BLOBFILEs).
  3. Recarrega o programa (load).
• Resultado: Como o texto MIL e as chaves do dicionário de pesos permanecem idênticos, o identificador interno do programa é preservado. Isso elimina a chamada ao compilador (ANECCompile), reduzindo o tempo de atualização de 4.200 ms para 494 ms (aceleração de 8.5x).

B. Treinamento Estável com Resume

O Orion resolveu três bugs críticos que causavam divergência para NaN no sistema anterior (ANEgpt):

1. Programas Obsoletos: Compilação feita antes do carregamento do checkpoint. Fix: Compilação deferida após o carregamento.
2. Overflow em FP16: O ANE opera nativamente em FP16. Fix: Clamp (limitação) de ativações antes de softmax e normalização.
3. Corrupção de Pesos: Dados de gradiente corrompidos ao escrever no BLOBFILE. Fix: Sanitização de gradientes (substituindo NaN/Inf por valores seguros).

• Resultado: Treinamento estável de 1.000 passos sem nenhum NaN, com resume de checkpoint funcionando perfeitamente.

C. LoRA como Entrada (Adapter-as-Input)

Para permitir a adaptação de modelos sem recompilar:

• As matrizes de baixo rank (LoRA) são passadas como entradas IOSurface dinâmicas, em vez de serem "assadas" nos pesos.
• Isso permite a troca "hot-swap" de adaptadores sem recompilação ou invalidação de cache.

---

4. Resultados e Desempenho

Os experimentos foram realizados em um Mac Studio com chip M4 Max.

• Inferência (GPT-2 124M):

  • Alcançou 170+ tokens/segundo no modo de decodificação (decode) no ANE.
  • Nota: A CPU (via Metal/cblas) foi ligeiramente mais rápida (283 tok/s) devido ao overhead de ida e volta do IOSurface (~2.3 ms) por despacho no ANE, mas o ANE oferece consumo de energia zero quando ocioso.
  • 100% de acordo com a base de referência da CPU em termos de sequência de tokens gerada.

• Treinamento (Stories110M em TinyStories):

  • v1.0 (Recompilação Total): ~85 minutos para 1.000 passos (com reinício de processo a cada passo).
  • v2.0 (Compilação Delta): 22 minutos para 1.000 passos em um único processo.
  • Aceleração Total: 3.8x no tempo de parede.
  • O tempo de recompilação caiu de 83.9% para 36.8% do tempo total do passo.
  • Zero ocorrências de NaN em 1.000 passos.

• Aceleração Específica:

  • Softmax com vocabulário grande (32k) é 33.8x mais rápido no ANE do que na CPU.
  • Camadas de classificação (forward) são 10.2x mais rápidas.

---

5. Significado e Impacto

• Desbloqueio de Hardware: O Orion demonstra que o ANE, presente em bilhões de dispositivos, é uma aceleração viável para treinamento e inferência de LLMs, não apenas para inferência simples.
• Sistema Completo: É a primeira implementação que combina caracterização de hardware, compilador, runtime estável e suporte a treinamento com checkpoint.
• Método de Engenharia: A técnica de "Compilação Delta" resolve o gargalo fundamental de modelos com pesos fixos na compilação, uma abordagem que pode ser aplicável a outros NPUs com comportamentos similares.
• Código Aberto: O projeto é lançado sob licença MIT, fornecendo a comunidade com ferramentas para explorar o ANE diretamente, incluindo conversores de pesos, tokenizadores nativos e benchmarks.

Em resumo, o Orion transforma o Apple Neural Engine de um "acelerador escuro" em uma plataforma programável e eficiente para o ciclo de vida completo de LLMs, superando limitações de software e hardware através de engenharia de sistemas inovadora.