Transformer de 25K parámetros en Commodore 64: Desafíos y Soluciones de Cuantificación y Ensamblador

La tesis central de este proyecto es demostrar que la arquitectura fundamental de los modelos Transformer, que impulsa los LLM modernos, puede ser implementada y ejecutada en hardware con recursos computacionales extremadamente limitados, como un Commodore 64 de 1 MHz. Esto aborda el problema fundamental de la eficiencia computacional y la portabilidad de modelos de inferencia, un desafío persistente en la historia de la computación. En un momento donde los modelos de IA crecen exponencialmente en tamaño y requisitos de hardware, este trabajo resalta la importancia de la optimización de bajo nivel y las técnicas de cuantificación para hacer la IA accesible y sostenible. No es solo un ejercicio de nostalgia, sino una exploración de los límites de la computación con restricciones severas, conectando las técnicas modernas de Machine Learning con los fundamentos de la programación de sistemas y la aritmética de punto fijo.

Históricamente, la computación ha avanzado desde máquinas con recursos mínimos hacia sistemas distribuidos masivos. Sin embargo, la necesidad de ejecutar algoritmos complejos en el 'edge' o en dispositivos de baja potencia persiste. Este proyecto es un ejemplo extremo de cómo los principios de la computación eficiente, como la reducción de precisión numérica y la optimización de código a nivel de ensamblador, son tan relevantes hoy como lo fueron en los albores de la informática. La capacidad de ejecutar un Transformer en un C64 subraya que la inteligencia artificial no es exclusiva de los superordenadores, sino que sus principios pueden ser destilados y adaptados a cualquier escala, siempre que se comprendan y manipulen las restricciones del hardware subyacente.

Arquitectura del Sistema

El sistema se compone de dos partes principales: el entorno de entrenamiento en Python y el motor de inferencia en ensamblador 6502/6510 para el Commodore 64. El modelo es un Transformer decoder-only de 2 capas, con 4 cabezas de atención (8 dimensiones por cabeza), embeddings de 32 dimensiones y 64 unidades ocultas en la FFN, totalizando ~25.000 parámetros cuantificados a INT8. La arquitectura de cada capa sigue el patrón estándar: RMSNorm → multi-head causal self-attention → residual → RMSNorm → ReLU MLP → residual, con una RMSNorm final y una proyección de salida seguida de argmax.

La clave de la implementación en el C64 reside en la aritmética de punto fijo Q8.8 (INT16) para todas las activaciones y la cuantificación INT8 para los pesos, con escalado por desplazamiento de potencia de 2 por tensor. Los sesgos son INT16 pre-escalados. Dado que el 6502 carece de instrucción de multiplicación, todas las operaciones matriciales se implementan mediante secuencias de desplazamiento y suma. La función Softmax utiliza una tabla de búsqueda (LUT) de 128 entradas para la función exponencial, con una normalización de scores de atención mediante un desplazamiento de 14 bits (>>14) para asegurar un rango dinámico adecuado. Esta normalización fue crítica para evitar que la atención se volviera uniforme. El motor de inferencia reside en la memoria del C64, con código y tablas de tokenizador en $0801-$20FF, pesos en $2100-$85A0 y buffers de activación en $8600-$9D00. El entrenamiento utiliza Quantization-Aware Training (QAT) con FakeQuantI8 y un 'straight-through gradient estimator', donde el desajuste deliberado entre el escalado continuo del entrenamiento y el escalado por potencia de 2 de la exportación actúa como ruido implícito, forzando al modelo a aprender pesos robustos a la cuantificación. El entrenamiento evalúa el 'forward pass' entero cada 500 épocas y guarda el mejor checkpoint basado en la precisión de argmax INT8.

class FakeQuantI8(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, scale):
        return (x / scale).round().clamp(-128, 127) * scale

    @staticmethod
    def backward(ctx, grad_output):
        return grad_output, None

def fq(x, scale):
    return FakeQuantI8.apply(x, scale)

# Dentro del bucle de entrenamiento:
# weights = fq(weights_float, scale_tensor)
# biases = fq(biases_float, scale_bias)

; Multiplicar A por 5 (ejemplo simplificado)
LDA operandA
ASL A       ; A = A * 2
STA temp
ASL A       ; A = A * 4
CLC
ADC temp    ; A = A * 4 + A * 1 = A * 5

Fundamentos Teóricos

Este proyecto se conecta directamente con los fundamentos de la teoría de la información y la computación con recursos limitados. La cuantificación de modelos, especialmente la cuantificación a enteros de baja precisión como INT8, tiene sus raíces en trabajos académicos sobre aritmética de punto fijo y representación numérica, esenciales para la computación embebida y de bajo consumo. El uso de tablas de búsqueda (LUT) para funciones complejas como la exponencial en Softmax es una técnica clásica en el diseño de hardware y software para optimizar el rendimiento en sistemas con recursos limitados, un principio explorado en la literatura de diseño de VLSI y microcontroladores desde hace décadas. La arquitectura Transformer en sí misma se basa en el paper "Attention Is All You Need" (Vaswani et al., 2017), que introdujo el mecanismo de auto-atención y la arquitectura que ha dominado el procesamiento de lenguaje natural. La adaptación de esta arquitectura a un entorno de 1 MHz y 64 KB de RAM es un testimonio de la flexibilidad de los principios algorítmicos cuando se combinan con una profunda comprensión de las restricciones del hardware. La implementación de RMSNorm, un componente clave en muchos Transformers modernos, también tiene sus raíces en trabajos recientes sobre normalización de capas en redes neuronales, como "Root Mean Square Layer Normalization" (Zhang & Sennrich, 2019), que busca mejorar la estabilidad del entrenamiento y la eficiencia computacional.