Escrito por Sarah Chieng, Head of DevX (@MilksAndMatcha
Una de las preguntas que más nos hacen es cómo desplegamos modelos grandes en los chips a escala de oblea (wafer-scale) de Cerebras. Una sola oblea Cerebras WSE-3T tiene 44 GB de memoria integrada en el chip, y los modelos frontera de billones de parámetros necesitan terabytes solo para sus pesos.
En este artículo, te explicaremos cómo dividimos un modelo frontera grande entre varias obleas WSE-3T, para ejecutar modelos de varios billones de parámetros a 1000-2000+ tokens por segundo.
Usaremos el modelo de código abierto Kimi K3, con ~3 billones de parámetros, como ejemplo ilustrativo.
Al terminar, entenderás:
- Cómo logramos que el modelo completo entre en la memoria del chip
- Cómo distribuimos el modelo entre varios chips
- Cómo manejamos una gran cantidad de solicitudes simultáneas
Desglosándolo con Kimi K3
Kimi K3 es un modelo de código abierto de 3 billones de parámetros con 69 capas de atención lineal y 24 capas densas.
Obtiene un impresionante 88.3% en Terminal-Bench 2.1 (frente al 88.8% de GPT-5.6 Sol) y un 91.2% en BrowseComp (frente al 90.4% de GPT-5.6 Sol). Estos resultados demuestran capacidades de programación y agentes de nivel frontera, superando incluso a los principales modelos propietarios en algunos benchmarks. Lee más sobre Kimi K3 aquí.
Para este ejemplo ilustrativo, usamos:
- Un tamaño promedio de solicitud asumido de 256k
- Cuantización de pesos de 4 bits (como en la implementación de Hugging Face
- Una caché KV de 8 bits
#1: Tamaño de cada capa
Para calcular el número total de obleas Cerebras WSE-3T que necesitamos para soportar Kimi K3, primero calculemos el tamaño de cada capa del modelo.
Kimi K3 tiene 3 billones de parámetros. Si los 3 billones de parámetros se almacenan a 4 bits cada uno, la carga útil bruta de parámetros es:
Dividir esa carga útil entre 93 capas nos da entonces:
#2: Número de obleas
Ahora que sabemos que cada capa necesita unos 16 GB, podemos calcular cuántas capas caben en cada oblea de 44 GB. La respuesta depende del tipo de capa:
- Las capas densas necesitan espacio adicional para una caché KV, porque cada capa densa (de atención completa) almacena las claves y los valores de cada token en cada solicitud activa. Por eso, colocamos una capa densa por oblea y dejamos la SRAM restante para su caché KV.
- Las capas de atención lineal no requieren espacio adicional para una caché KV. En lugar de almacenar claves y valores para cada token, una capa de atención lineal mantiene un estado pequeño y de tamaño fijo por solicitud, sin importar la longitud del contexto. Eso deja espacio para dos capas lineales por oblea (2 × 16 GB = 32 GB).
Por lo tanto, con las 69 capas lineales y 24 capas densas de Kimi K3, obtenemos:
Esto nos da un total de 35 + 24 = 59 obleas.
#3: Conectando las obleas entre sí
Cerebras ejecuta los modelos en un pipeline. Después de que el compilador de Cerebras distribuye las capas del modelo entre distintas regiones de las obleas, el entorno de ejecución traza la ruta para que los tokens fluyan entre ellas. Para lograr una latencia ultrarrápida, solo es necesario pasar las activaciones entre las obleas, y esto puede hacerse mediante RDMA sobre Ethernet. Esto toma aproximadamente 2 microsegundos por salto.
Es importante destacar que, dado que cada oblea de Cerebras tiene 44 GB de SRAM integrada y cada capa ocupa solo 16 GB, normalmente podemos colocar capas enteras del modelo directamente en el chip.
Esto significa que, en la mayoría de los casos, los pesos del modelo nunca tienen que moverse y las capas no necesitan dividirse entre varias obleas. Solo las activaciones pasan de una oblea a la siguiente.
#4: Cómo atendemos 1 solicitud
Ahora que tenemos las obleas conectadas y el modelo distribuido, podemos pensar en cada oblea como una etapa en una línea de ensamblaje: los pesos del modelo de cada etapa permanecen en su lugar, y los resultados intermedios de la solicitud avanzan por las etapas de las capas.
Por ejemplo, la Oblea 1 ejecuta las capas 1–2 del modelo y envía sus activaciones a la oblea 2. La Oblea 2 ejecuta las capas 3–4 y pasa sus resultados a la Oblea 3, y así sucesivamente hasta que la última oblea produce el siguiente token.
#5: Cómo atendemos muchas solicitudes
Aquí surge un malentendido común. Si un token tiene que recorrer cada oblea en orden, ¿acaso cada oblea no queda inactiva mientras espera su turno? En efecto, ¿no estaría funcionando el sistema con un batch size de uno?
¡No! Si solo enviáramos una solicitud a la vez por la línea y esperáramos a que terminara antes de iniciar la siguiente, estaríamos usando el pipeline de forma ineficiente. En cambio, mantenemos todo el pipeline lleno con solicitudes distintas. Como el modelo ya está dividido en etapas de pipeline, cada etapa puede trabajar en una solicitud diferente al mismo tiempo.
Mientras la etapa 3 procesa un token para la solicitud A, la etapa 2 procesa uno para la solicitud B y la etapa 1 uno para la solicitud C. En nuestro ejemplo, cada etapa es una capa. En producción, Cerebras usa etapas mucho más detalladas, y muchas más de ellas.
Cada etapa también puede procesar un micro-batch, un grupo pequeño de solicitudes, en lugar de una sola solicitud. Así, la cantidad de solicitudes en tránsito es aproximadamente el número de etapas multiplicado por el tamaño del micro-batch.
Para entender el balance: los micro-batches más pequeños hacen que cada solicitud sea aún más rápida, mientras que los micro-batches más grandes aumentan el rendimiento total.
#6: ¿Cuántas solicitudes podemos atender?
La cantidad de memoria del chip asignada a la caché KV determina el número máximo de solicitudes simultáneas. Cada solicitud en tránsito necesita su propia caché KV en cada oblea de capa densa, por lo que el tamaño de la memoria del chip establece el límite de solicitudes concurrentes. Como explicamos arriba, solo necesitamos calcular los requisitos de caché KV para las capas densas.
Para cada oblea de 44 GB que ya tiene una capa densa, podemos calcular el espacio restante para la caché KV con:
Almacenamiento total de la oblea - tamaño de una capa densa. En este caso:
Calculando el n.º de solicitudes simultáneas
Para encontrar el número de solicitudes simultáneas soportadas para cualquier distribución particular del modelo, primero calculamos cuánta caché ocupa una sola solicitud de 256k:
Luego, para calcular el número total de solicitudes simultáneas que esta configuración soportaría, dividimos:
El número total de solicitudes simultáneas representa la cantidad de solicitudes que se mueven por el pipeline al mismo tiempo en cualquier momento. A diferencia de un batch en GPU, esas solicitudes no esperan a que las demás terminen; fluyen por el pipeline de manera continua.
Por supuesto, ni la concurrencia ni el tamaño del batch te dicen exactamente qué tan eficiente es un sistema de inferencia. Al evaluar el rendimiento de un sistema de inferencia, la métrica más significativa es el rendimiento total de tokens por unidad de energía (por ejemplo, tokens por segundo por megavatio), cuando se mide a la misma velocidad de generación de salida.
Cómo escala esto más allá de los 3 billones de parámetros
El mismo enfoque aplica para modelos más grandes: calcula la memoria necesaria para los pesos y el estado de las solicitudes, distribuye el modelo entre las obleas y conéctalas en un pipeline. Cada oblea adicional aporta capacidad de cómputo y memoria integrada, lo que nos permite manejar modelos que superan la capacidad de cualquier chip individual.
¡Gracias por aprender con nosotros! ¿Te interesa desplegar un modelo en Cerebras o desarrollar con nuestra nube de inferencia?
- Comienza con nuestra API gratuita: cloud.cerebras.ai
- Contacta a nuestro equipo de ventas: sales@cerebras.ai
- Lee más artículos en cerebras.ai
- Lee el artículo original aquí.
Diseños de @halleychangg. Gracias a Atharva Talpade, Zhenwei Gao, Angela Yeung, Jessica Liu, Joyce Er, Emma Call y Rohan Gupta por su revisión y comentarios.
