NeuronFabric: A Software Reference Architecture for On-Chip Transformer Training with Local Adam

Resumen Técnico: NeuronFabric

Planteamiento del Problema

Los aceleradores neuronales actualmente documentados públicamente generalmente separan la fase de cómputo de la fase de actualización del estado del optimizador. Los chips dedicados únicamente a la inferencia (p. ej., Groq LPU, Apple ANE, IBM NorthPole) carecen de computación de gradientes en el chip, mientras que los aceleradores de entrenamiento (p. ej., Cerebras, Tenstorrent) típicamente delegan las actualizaciones de pesos a una CPU host o a una GPU externa. Esta separación arquitectónica crea cuellos de botella significativos:

• Tráfico de Pesos: La inferencia requiere cargar todos los pesos desde la memoria externa (HBM) para cada pasada hacia adelante (forward pass).
• Latencia de Entrenamiento: El chip que computa los gradientes no los aplica; las actualizaciones se orquestan por separado, lo que a menudo requiere el movimiento de datos fuera del chip.
• Limitaciones de Aprendizaje Continuo: Los chips de inferencia no pueden actualizar los pesos a partir de nuevos datos sin recargar el modelo completo desde un clúster externo, un proceso que toma de minutos a horas.

El artículo postula que un chip de silicio donde cada neurona posea sus propios pesos y ejecute el optimizador Adam localmente podría minimizar el tráfico fuera del chip y permitir un verdadero entrenamiento en el chip. Sin embargo, antes de comprometerse con el silicio, la validez matemática y la corrección numérica de tal arquitectura de "Adam local" deben establecerse.

Metodología

Los autores desarrollaron NeuronFabric, una arquitectura de referencia de software completa implementada en C# sin marcos de trabajo (frameworks) externos de aprendizaje automático. El sistema está diseñado de modo que cada abstracción de software se mapee directamente a un bloque de hardware (p. ej., NeuronCore a BRAM, Forward() a cadenas de DSP48).

Arquitectura Central:

• Adam Local: Cada NeuronCore almacena su propio vector de pesos $w$ y sesgo $b$. Durante la pasada hacia atrás (backward pass), computa gradientes y actualiza los pesos in-place utilizando el algoritmo Adam. No existe un bus de pesos entre los núcleos.
• Configuración del Modelo: Un transformador autorregresivo de 334K parámetros (Pre-LN) entrenado con el corpus de Shakespeare.
  • Dimensiones: Embedding $d=88$, 4 cabezales, Feedforward $f=264$, 4 capas.
  • Vocabulario: 256 (nivel de byte).
  • Técnicas: Vinculación de pesos (weight tying; la proyección de salida reutiliza la matriz de embedding) para ahorrar parámetros.

Esquema de Precisión BF16W:
Para abordar las estrictas restricciones de SRAM en el entrenamiento en FPGA, los autores introdujeron BF16W (Brain Floating Point 16 Weights).

• Almacenamiento: Los pesos se almacenan como BF16 (2 bytes).
• Computación: Los pesos se convierten a FP32 para la computación de la pasada hacia adelante/atrás y las actualizaciones de los momentos de Adam.
• Momentos: Los primeros ($m$) y segundos ($v$) momentos de Adam permanecen en FP32 (4 bytes cada uno) para asegurar la estabilidad.
• Actualización: Después del paso de Adam, el peso actualizado en FP32 se redondea de nuevo a BF16 y se escribe en el almacenamiento.

Restricción del Presupuesto de Vocabulario:
El artículo identifica una restricción crítica para modelos de pocos parámetros: la tabla de embeddings ($|V| \times d$) consume una parte desproporcionada del presupuesto de parámetros ("impuesto de vocabulario"), dejando capacidad insuficiente para que las capas del transformador aprendan patrones estructurales. Esto se cuantificó a través de tres dominios, mostrando que cuando la capacidad de razonamiento restante ($P_{reason}$) cae por debajo de ~20K, la generación de texto coherente falla.

Contribuciones Clave

1. Implementación de Referencia de Software: Una implementación completa en C# de un transformador sin dependencias de frameworks, con retropropagación completa y actualizaciones de Adam local, que ejecuta entrenamiento, evaluación y guardado de puntos de control (checkpointing) en un solo binario.
2. Validación de BF16W: Demostración de que almacenar los pesos en BF16 mientras se mantienen los momentos en FP32 reduce el uso de SRAM de 12 bytes/parámetro a 10 bytes/parámetro con una penalización de convergencia insignificante (+0.020 en la pérdida de validación en Shakespeare).
3. Cuantificación del Presupuesto de Vocabulario: Evidencia empírica que muestra que, para un hardware de presupuesto fijo, la tabla de embeddings domina el modelo, lo que requiere vocabularios pequeños o arquitecturas especializadas de múltiples chips (p. ej., chips de embedding separados).
4. Objetivo de Viabilidad en FPGA: Una configuración concreta (modelo de 334K BF16W) que cabe enteramente dentro de los 4.0 MB de BRAM en el chip de una FPGA Xilinx ZCU102, dejando 660 KB de margen para buffers de activación, sin requerir acceso a DDR durante el entrenamiento.

Resultados

El sistema fue validado en una máquina con Windows 11 (CPU AMD Ryzen 9 9900X, GPU NVIDIA RTX 4090) utilizando el corpus de Shakespeare (80K muestras).

• Correctitud: Las comprobaciones numéricas de gradientes (diferencias finitas vs. gradientes analíticos) pasaron para todos los componentes, confirmando la corrección matemática de la implementación de Adam local.
• Convergencia:
  • GPU FP32 (Oráculo): Logró una pérdida de validación de 1.5224.
  • CPU BF16W: Logró una pérdida de validación de 1.5426.
  • La brecha de +0.020 se atribuye a la diferencia de precisión entre el almacenamiento de pesos en BF16 y la computación en FP32.
• Calidad Generativa: El modelo produjo texto coherente a nivel de caracteres con estructuras de diálogo de Shakespeare reconocibles y patrones de nombres de personajes, confirmando que el bucle de entrenamiento converge a una solución funcional.
• Eficiencia de SRAM:
  • Adam FP32 (12 bytes/parámetro): 4.00 MB totales (llena exactamente la BRAM de la ZCU102, dejando cero espacio para activaciones).
  • Adam BF16W (10 bytes/parámetro): 3.34 MB totales (deja 660 KB libres para activaciones y estado de control).

Significado y Reivindicaciones

El artículo establece explícitamente su alcance modesto: NeuronFabric es un prototipo de investigación, no un acelerador de producción. No contiene mediciones de FPGA ni resultados de cierre de tiempo (timing closure).

La importancia principal del artículo radica en:

1. Establecer el Estado de la Técnica: Sirve como una divulgación arquitectónica pública y una referencia de software para establecer la arquitectura de "actualización de Adam local" como estado de la técnica para futuras implementaciones en FPGA/ASIC.
2. Prueba de Concepto: Demuestra que las matemáticas funcionan y que los números son correctos para un transformador que entrena enteramente en el chip con el estado del optimizador local, sin una CPU host en el bucle de actualización.
3. Definir el Camino a Seguir: Establece el modelo de Shakespeare de 334K BF16W como un objetivo falsable para una futura implementación en FPGA. Si se puede lograr la misma curva de pérdida en la FPGA sin intervención del host, la viabilidad del entrenamiento de transformadores íntegramente en el chip se convierte en un hecho medible en lugar de una proyección de software.

Los autores concluyen que el siguiente paso inmediato es implementar este bucle de entrenamiento específico en la FPGA ZCU102 para verificar que los resultados de software se traduzcan al hardware, validando así el paradigma de entrenamiento "sin CPU host".