El problema fundamental que aborda este enfoque es la limitación de memoria y capacidad de cómputo de las arquitecturas de hardware convencionales (ej. GPUs) para modelos de lenguaje de escala de terabytes. A medida que los modelos de IA crecen exponencialmente en tamaño, superando la memoria de un solo dispositivo, surge la necesidad de estrategias de distribución eficientes que minimicen la latencia de comunicación y maximicen el throughput. La solución propuesta por Cerebras, utilizando sus Wafer-Scale Engines (WSE), se centra en una arquitectura de memoria distribuida y un pipeline de ejecución que permite mantener los pesos del modelo en SRAM on-chip y mover solo las activaciones entre wafers, mitigando los cuellos de botella de ancho de banda de memoria externa.
La relevancia actual de este problema es crítica, ya que la carrera por modelos de lenguaje más grandes y capaces continúa. La capacidad de ejecutar estos modelos de manera eficiente, tanto en términos de latencia como de throughput, es un diferenciador clave para la inferencia a escala de producción. Este enfoque se distingue por su énfasis en la memoria on-chip y la minimización del movimiento de datos, un principio fundamental en el diseño de sistemas de alto rendimiento.
Arquitectura del Sistema
La arquitectura se basa en la distribución de las capas de un modelo de lenguaje (ej. Kimi K3 de 3 billones de parámetros) a través de múltiples wafers Cerebras WSE-3T. Cada WSE-3T posee 44 GB de SRAM on-chip. Las capas del modelo, con un tamaño aproximado de 16 GB cada una (tras cuantificación a 4 bits), se asignan a los wafers. Las capas densas (full-attention) requieren espacio adicional para un KV cache dinámico, por lo que se asigna una capa densa por wafer. Las capas de atención lineal, que mantienen un estado de tamaño fijo por solicitud, permiten alojar dos capas por wafer.
La interconexión entre wafers se realiza mediante RDMA sobre Ethernet, con una latencia de aproximadamente 2 microsegundos por 'hop'. La ejecución del modelo sigue un patrón de pipeline: cada wafer actúa como una etapa, procesando una o dos capas del modelo y pasando las activaciones resultantes al siguiente wafer. Los pesos del modelo permanecen estáticos en la SRAM de cada wafer, eliminando la necesidad de cargar y descargar datos de memoria externa. Para manejar múltiples solicitudes concurrentes, el pipeline se mantiene lleno, con diferentes solicitudes siendo procesadas simultáneamente en distintas etapas. Cada etapa puede procesar un 'micro-batch' de solicitudes, permitiendo un trade-off entre latencia por solicitud (micro-batches pequeños) y throughput total (micro-batches grandes). La capacidad de KV cache en los wafers con capas densas determina el número máximo de solicitudes concurrentes.
Flujo de Inferencias en Pipeline
- 1 Wafer 1 (Capas 1-2) Procesa activaciones de la Solicitud C, envía resultados a Wafer 2.
- 2 Wafer 2 (Capas 3-4) Procesa activaciones de la Solicitud B, envía resultados a Wafer 3.
- 3 Wafer 3 (Capas 5-6) Procesa activaciones de la Solicitud A, envía resultados a Wafer 4.
- 4 Wafer N (Últimas Capas) Genera el siguiente token para su solicitud, envía a salida.
| Capa | Tecnología | Justificación |
|---|---|---|
| compute | Cerebras WSE-3T | Unidad de procesamiento principal con 44 GB de SRAM on-chip, diseñada para cómputo de IA a escala de wafer. vs GPU Clusters, TPU Clusters |
| storage | SRAM on-chip | Almacena los pesos del modelo y el KV cache para cada capa, eliminando accesos a memoria externa de alta latencia. vs HBM (High Bandwidth Memory), DDR DRAM 44 GB por wafer |
| networking | RDMA over Ethernet | Proporciona comunicación de baja latencia (2 µs por hop) para el paso de activaciones entre wafers en el pipeline. vs PCIe, Infiniband |
Trade-offs
Ganancias
- ▲ Throughput de tokens
- △ Latencia de inferencia (por solicitud)
- ▲ Eficiencia energética (tokens/segundo/megawatt)
Costes
- ▲ Complejidad de orquestación del modelo
- ▲ Costo inicial de hardware
Fundamentos Teóricos
Este enfoque se conecta con principios fundamentales de la computación paralela y distribuida, particularmente el concepto de 'pipeline parallelism' o 'model parallelism', que ha sido estudiado extensamente en la literatura de computación de alto rendimiento. La idea de dividir una tarea computacional en etapas secuenciales y procesar diferentes elementos de datos en cada etapa simultáneamente se remonta a los diseños de procesadores RISC y arquitecturas vectoriales. En el contexto de redes neuronales profundas, el paper 'GPipe: Efficient Training of Giant Models using Pipeline Parallelism' (Huang et al., 2019) formalizó el uso de esta técnica para el entrenamiento de modelos grandes, aunque aquí se aplica a la inferencia.
La optimización del uso de memoria on-chip y la minimización del movimiento de datos entre niveles de la jerarquía de memoria es un tema recurrente en la arquitectura de computadoras, relacionado con el 'memory wall' y el diseño de sistemas 'data-centric'. La elección de SRAM on-chip para los pesos del modelo y el KV cache refleja la aplicación de estos principios para reducir drásticamente la latencia y el consumo de energía asociados con el acceso a DRAM externa, un cuello de botella bien documentado en la computación moderna.