Treinando Redes Neurais Diretamente no Apple Neural Engine Sem CoreML e GPU

Todo processador Apple Silicon contém um bloco dedicado para operações de redes neurais — o Apple Neural Engine (ANE). O fabricante alega dezenas de teraflops de desempenho, mas as ferramentas oficiais limitam seu uso apenas à inferência via CoreML. O treinamento de modelos é oferecido tanto na CPU quanto no chip gráfico.

Um desenvolvedor com o apelido maderix decidiu testar se o hardware do ANE é fisicamente incapaz de fazer backpropagation, ou se o problema é apenas devido às limitações de software da Apple. Em um fim de semana, ele dissecou os frameworks privados do macOS e conseguiu fazer o ANE executar não apenas o forward pass, mas também o backward pass para transformers.

O projeto reuniu mais de 7 mil estrelas no GitHub. Vamos analisar como esse hack funciona, quais truques tiveram que ser aplicados e como é o desempenho real.

Por Que Investigar Frameworks Fechados

A stack oficial como o CoreML cria a impressão de que o ANE é uma caixa preta fechada. Você fornece um modelo pronto, ele retorna o resultado. Se você precisa treinar até um pequeno modelo diretamente no cliente, era necessário recorrer ao Metal ou ao framework MLX, carregando a GPU.

O autor do repositório ANE demonstrou que o chip é perfeitamente capaz de executar grafos computacionais arbitrários. Para isso, ele fez engenharia reversa das bibliotecas privadas _ANEClient e _ANECompiler, além da linguagem interna de descrição de modelos MIL (Model Intermediate Language).

O resultado foi interessante: treinamento completo de transformers sem uma única linha de CoreML ou Metal. O texto do modelo é montado diretamente na RAM, compilado em tempo real e enviado ao processador neural.

Limitações e Realidade Rigorosa

Antes de correr para reescrever seus scripts de treinamento no Mac, vale a pena olhar os números honestos. O próprio autor adverte desde o início: isso é um experimento acadêmico, não uma biblioteca pronta para produção.

Alcançar 100% de utilização dos recursos do ANE ainda não foi possível. A utilização real do chip é de cerca de 5–9% do pico. Limitações de software e hardware se fazem notar:

• Algumas operações matemáticas não são suportadas pelo chip na forma necessária e recorrem à CPU.
• O backward pass para pesos (dW) ainda precisa ser calculado pelo processador.
• O compilador do ANE tem vazamentos de memória, então após cerca de cem iterações, são necessárias soluções alternativas.

No entanto, mesmo nesse modo, o projeto entrega uma velocidade decente em arquiteturas básicas.

Como o Pipeline de Treinamento Funciona

A arquitetura do projeto depende de distribuir tarefas entre ANE e CPU. O acelerador de rede neural lida com as multiplicações de matrizes mais pesadas, enquanto o processador cuida da lógica circundante e do acúmulo de gradientes.

O forward pass e o cálculo de gradiente de entrada (dx) são totalmente executados no ANE. Os gradientes de pesos (dW) são calculados pela CPU através das bibliotecas otimizadas Accelerate e cblas_sgemm. O otimizador Adam e a camada RMSNorm também são executados no processador.

Para evitar recompilar o grafo do modelo a cada passo quando os pesos mudam, o autor aplicou um truque: pesos e ativações são compactados em um único tensor através das dimensões espaciais, e dentro do kernel MIL eles são simplesmente separados novamente.

A memória IOSurface é usada para troca de dados entre CPU e ANE. Isso torna possível transferir tensores sem cópia desnecessária entre espaços de endereço. Os dados são empacotados no formato específico do ANE [1, C, 1, S], onde os canais vêm primeiro. Essa abordagem eliminou a sobrecarga de transposição de matrizes.

Problemas Não Óbvios e Soluções Alternativas

Muitas armadilhas do hardware Apple surgiram durante a engenharia reversa.

Primeiro, a operação SDPA (Scaled Dot-Product Attention) no ANE ignora a máscara causal attn_mask no nível de hardware. O mecanismo de atenção teve que ser dividido em três estágios: multiplicação de Q e K no ANE, mascaramento com softmax na CPU, e multiplicação final por V de volta no ANE.

Segundo, o compilador embutido _ANECompiler contém um vazamento de memória. Após aproximadamente 119 compilações, o processo trava devido ao esgotamento de recursos. O autor resolveu o problema radicalmente: quando o contador se aproxima do limite, o programa salva um checkpoint e faz exec() — reinicia a si mesmo com preservação de estado.

Terceiro, cálculos em FP16 durante o backward pass rapidamente levam a underflow, fazendo os gradientes se transformarem em zeros. O problema foi resolvido escalando a loss com o coeficiente 256 * NLAYERS.

Desempenho no M4

No chip Apple M4, os resultados ficaram bastante ilustrativos. Os testes foram conduzidos em duas arquiteturas:

Para o modelo Stories110M com 109 milhões de parâmetros (12 camadas, Multi-Head Attention clássico), o tempo para um passo de treinamento foi de 91 milissegundos.

O maior Qwen3-0.6B com 596 milhões de parâmetros e Grouped-Query Attention processa um passo em 412 milissegundos.

O autor também testou quantização INT8 W8A8. Usar pesos e ativações de 8 bits reduz a carga na memória L2 SRAM do chip e aumenta a vazão de 18.6 TOPS para 35.1 TOPS no M4 — um speedup de quase 1.88x comparado ao FP16.

Como Executar o Projeto

O projeto não requer dependências externas como PyTorch ou Conda. Tudo que você precisa é do macOS 15 mais recente em uma máquina com Apple Silicon e o compilador Clang. As APIs privadas são carregadas em tempo de execução via objc_msgSend.

Para construir o pipeline dinâmico, basta navegar até a pasta do projeto e executar o comando make:

cd training/training_dynamic
make MODEL=stories110m
./train --scratch

Se você quiser testar a quantização INT8 ou medir os TOPS de pico do seu chip, existem benchmarks separados na pasta raiz do repositório.

Considerações Finais

O projeto maderix é um exemplo de pesquisa de qualidade "sob o capô" do hardware Apple. Ele demonstra que as limitações do ANE estão unicamente no plano de software e na natureza fechada do ecossistema.

Usar o repositório para treinar grandes modelos de linguagem em produção atualmente é inútil — GPU e o framework MLX existem para isso. Mas se você está estudando como aceleradores neurais funcionam, escrevendo seus próprios compiladores para Edge AI, ou quer entender como trabalhar com APIs privadas do macOS diretamente de C e Objective-C, esse código será um excelente recurso de aprendizado.