Entrenamiento de Redes Neuronales Directamente en Apple Neural Engine Sin CoreML ni GPU

Cada procesador Apple Silicon contiene un bloque dedicado para operaciones de redes neuronales — el Apple Neural Engine (ANE). El fabricante afirma decenas de teraflops de rendimiento, pero las herramientas oficiales limitan su uso a inferencia únicamente a través de CoreML. El entrenamiento de modelos se ofrece ya sea en CPU o en el chip gráfico.

Un desarrollador con el apodo maderix decidió probar si el hardware ANE es físicamente incapaz de hacer backpropagation, o si el problema se debe únicamente a las limitaciones de software de Apple. Durante un fin de semana, disectó los frameworks privados de macOS e hizo que ANE ejecutara no solo el paso hacia adelante sino también el paso hacia atrás para transformers.

El proyecto reunió más de 7 mil estrellas en GitHub. Desglosemos cómo funciona este hack, qué trucos se tuvieron que aplicar y cómo se ve el rendimiento real.

Por Qué Profundizar en Frameworks Cerrados

El stack oficial como CoreML crea la impresión de que ANE es una caja negra cerrada. Le das un modelo listo, te devuelve el resultado. Si necesitas entrenar incluso un modelo pequeño directamente en el cliente, tenías que recurrir a Metal o al framework MLX, cargando la GPU.

El autor del repositorio ANE demostró que el chip es perfectamente capaz de ejecutar grafos computacionales arbitrarios. Para hacer esto, reverse-engineereó las librerías privadas _ANEClient y _ANECompiler, así como el lenguaje interno de descripción de modelos MIL (Model Intermediate Language).

El resultado fue interesante: entrenamiento completamente funcional de transformers sin una sola línea de CoreML o Metal. El texto del modelo se ensambla directamente en RAM, se compila sobre la marcha y se envía al procesador neural.

Limitaciones y Realidad Cruda

Antes de apresurarte a reescribir tus scripts de entrenamiento en Mac, vale la pena ver números honestos. El propio autor advierte de entrada: esto es un experimento académico, no una librería lista para producción.

No ha sido posible lograr el 100% de utilización de los recursos ANE. La utilización real del chip es de aproximadamente 5–9% del pico. Las limitaciones de software y hardware se hacen notar:

• Algunas operaciones matemáticas no son compatibles con el chip en la forma requerida y recurren a la CPU.
• El paso hacia atrás para los pesos (dW) todavía tiene que ser calculado por el procesador.
• El compilador de ANE tiene memory leaks, por lo que después de unas cien iteraciones se necesitan workarounds.

Sin embargo, incluso en este modo, el proyecto ofrece una velocidad decente en arquitecturas básicas.

Cómo Funciona el Pipeline de Entrenamiento

La arquitectura del proyecto se basa en distribuir tareas entre ANE y CPU. El acelerador de redes neuronales maneja las multiplicaciones de matrices más pesadas, mientras que el procesador se encarga de la lógica circundante y la acumulación de gradientes.

El paso hacia adelante y el cálculo del gradiente de entrada (dx) se ejecutan completamente en ANE. Los gradientes de los pesos (dW) se calculan por CPU a través de las librerías optimizadas Accelerate y cblas_sgemm. El optimizador Adam y la capa RMSNorm también se ejecutan en el procesador.

Para evitar recompilar el grafo del modelo en cada paso cuando cambian los pesos, el autor aplicó un truco: los pesos y las activaciones se empaquetan en un solo tensor a través de las dimensiones espaciales, y dentro del kernel MIL simplemente se separan de nuevo.

Se usa memoria IOSurface para el intercambio de datos entre CPU y ANE. Esto hace posible transferir tensores sin copia innecesaria entre espacios de direcciones. Los datos se empaquetan en el formato específico de ANE [1, C, 1, S], donde los canales van primero. Este enfoque eliminó la sobrecarga de la transposición de matrices.

Problemas No Obvios y Workarounds

Muchas trampas del hardware de Apple surgieron durante el reverse engineering.

Primero, la operación SDPA (Scaled Dot-Product Attention) en ANE ignora la máscara causal attn_mask a nivel de hardware. El mecanismo de atención tuvo que dividirse en tres etapas: multiplicar Q y K en ANE, aplicar máscara con softmax en CPU, y la multiplicación final por V de vuelta en ANE.

Segundo, el compilador integrado _ANECompiler contiene un memory leak. Después de aproximadamente 119 compilaciones, el proceso falla debido al agotamiento de recursos. El autor resolvió el problema radicalmente: cuando el contador se acerca al límite, el programa guarda un checkpoint y hace exec() — se reinicia a sí mismo con preservación del estado.

Tercero, los cálculos en FP16 durante el paso hacia atrás rápidamente conducen a underflow, causando que los gradientes se conviertan en ceros. El problema se resolvió escalando la pérdida con el coeficiente 256 * NLAYERS.

Rendimiento en M4

En el chip Apple M4, los resultados resultaron bastante ilustrativos. Las pruebas se realizaron en dos arquitecturas:

Para el modelo Stories110M con 109 millones de parámetros (12 capas, Multi-Head Attention clásico), el tiempo para un paso de entrenamiento fue de 91 milisegundos.

El más grande Qwen3-0.6B con 596 millones de parámetros y Grouped-Query Attention procesa un paso en 412 milisegundos.

El autor también probó cuantización INT8 W8A8. Usar pesos y activaciones de 8 bits reduce la carga en la memoria SRAM L2 del chip y aumenta el throughput de 18.6 TOPS a 35.1 TOPS en M4 — una aceleración de casi 1.88x comparada con FP16.

Cómo Ejecutar el Proyecto

El proyecto no requiere dependencias externas como PyTorch o Conda. Todo lo que necesitas es la última versión de macOS 15 en una máquina con Apple Silicon y el compilador Clang. Las APIs privadas se cargan en tiempo de ejecución a través de objc_msgSend.

Para construir el pipeline dinámico, simplemente navega a la carpeta del proyecto y ejecuta el comando make:

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

Si quieres probar la cuantización INT8 o hacer benchmark del pico TOPS de tu chip, hay benchmarks separados en la carpeta raíz del repositorio.

Reflexiones Finales

El proyecto de maderix es un ejemplo de investigación de calidad "bajo el capó" del hardware de Apple. Demuestra que las limitaciones de ANE residen únicamente en el plano del software y la naturaleza cerrada del ecosistema.

Usar el repositorio para entrenar modelos de lenguaje grandes en producción actualmente es inútil — GPU y el framework MLX existen para eso. Pero si estás estudiando cómo funcionan los aceleradores neurales, escribiendo tus propios compiladores para Edge AI, o quieres entender cómo trabajar con las APIs privadas de macOS directamente desde C y Objective-C, este código será un excelente recurso de aprendizaje.