Paralelismo de modelos: cómo Kimi K3 reparte billones de parámetros

⏱️ Lectura: 15 min

Kimi K3 tiene 3 billones de parámetros. Guardados en 4 bits pesan 1,5 terabytes, mucho más que los 44 GB de memoria que trae un solo chip wafer-scale de Cerebras. La única forma de servir un modelo así es el paralelismo de modelos: partirlo entre decenas de chips a la vez.

📑 En este artículo

1. TL;DR
2. ¿Qué es el paralelismo de modelos?
3. Por qué importa
4. Cómo funciona el sharding de pesos y la atención entre chips
   1. Pipeline: cada chip ejecuta un tramo de capas
   2. Tensor: partir una misma capa en pedazos
   3. Expertos: solo activar una fracción del modelo
5. El caso Kimi K3 en los chips Cerebras WSE-3T
6. Cómo empezar: probar el sharding de pesos con Hugging Face Accelerate
7. Casos de uso reales
8. Errores comunes y buenas prácticas
9. Comparativa con alternativas
10. Profundizando
11. Preguntas frecuentes
    1. ¿El paralelismo de modelos reemplaza al paralelismo de datos?
    2. ¿Por qué Kimi K3 necesita 59 wafers Cerebras WSE-3T?
    3. ¿device_map=»auto» de Hugging Face hace sharding de pesos real?
    4. ¿El paralelismo por pipeline hace que un chip espere a que termine la petición anterior?
    5. ¿Qué diferencia hay entre paralelismo de tensor y paralelismo de expertos?
12. Referencias
    1. 📚 Artículos relacionados

No es exclusivo de Cerebras: los mismos principios corren en clústeres de GPU con Megatron-LM, DeepSpeed o vLLM. Esta guía explica cómo se reparten los pesos y la atención entre chips, usando Kimi K3 sobre wafers Cerebras como caso concreto, no como anuncio de producto.

TL;DR

• Kimi K3 reparte sus 93 capas (69 de atención lineal y 24 densas) entre 59 wafers Cerebras WSE-3T.
• Cada capa de Kimi K3 pesa unos 16 GB en 4 bits, así que ningún chip de 44 GB aguanta el modelo completo.
• Las capas densas reservan memoria para el cache KV; las lineales no, por eso caben dos por wafer en vez de una.
• device_map=»auto» de Hugging Face aplica el mismo principio de sharding en una GPU doméstica, sin clúster wafer-scale.
• Un pipeline de chips mantiene varias peticiones en vuelo a la vez, así ningún wafer queda ocioso esperando turno.

¿Qué es el paralelismo de modelos?

El paralelismo de modelos es el conjunto de técnicas que divide los pesos y los cálculos de una red neuronal entre varios chips cuando el modelo no entra en la memoria de un solo acelerador, repartiendo capas, porciones de cada matriz o grupos de expertos según la arquitectura del hardware disponible.

Existen tres variantes principales, y los sistemas grandes casi siempre combinan más de una: paralelismo por pipeline (cada chip ejecuta un grupo de capas distintas), paralelismo de tensor (cada chip ejecuta una porción de la misma capa) y paralelismo de expertos (cada chip guarda un subconjunto de redes especializadas en modelos tipo Mixture of Experts). Kimi K3 sobre Cerebras usa sobre todo la primera.

Por qué importa

El paralelismo de modelos no es una optimización opcional, sino una condición necesaria para que el modelo cargue. Un modelo de 3 billones de parámetros en 4 bits pesa 1,5 TB, y la GPU más grande del mercado trae decenas de gigabytes de memoria, no terabytes. Sin repartir el modelo, simplemente no entra en el chip.

Lo mismo pasa a menor escala. Un modelo de 70.000 millones de parámetros en precisión completa ya pesa unos 280 GB, más que la memoria combinada de varias GPUs de consumo. Por eso cualquier equipo que sirva un modelo grande, sea con wafers Cerebras o con un clúster de GPU, termina resolviendo el mismo problema de reparto.

Cómo funciona el sharding de pesos y la atención entre chips

Antes de repartir nada, hay que saber cuánto pesa cada pieza. El compilador de Cerebras mapea cada capa del modelo a una región física de wafers, y el runtime decide por dónde viajan los tokens entre esas regiones. El siguiente diagrama resume cómo quedan las 93 capas de Kimi K3 repartidas en wafers de 44 GB:

flowchart TD
    A["69 capas de atención lineal"] --> B["2 capas lineales por wafer (32 GB)"]
    C["24 capas densas"] --> D["1 capa densa + cache KV por wafer (44 GB)"]
    B --> E["35 wafers para capas lineales"]
    D --> F["24 wafers para capas densas"]
    E --> G["59 wafers WSE-3T en total"]
    F --> G

Pipeline: cada chip ejecuta un tramo de capas

En paralelismo por pipeline, cada chip guarda un grupo de capas consecutivas y actúa como una estación de una línea de ensamblaje. El wafer 1 ejecuta las capas 1 y 2 y manda sus activaciones al wafer 2, que ejecuta las capas 3 y 4 y las pasa al wafer 3, y así hasta la última etapa. Solo las activaciones cruzan de un wafer a otro, nunca los pesos, porque cada capa ya entera cabe en la memoria del wafer que la hospeda.

En el caso de Cerebras, esas activaciones viajan por RDMA sobre Ethernet, con una latencia de unos 2 microsegundos por salto. El siguiente diagrama muestra cómo se llena el pipeline con varias peticiones en simultáneo:

sequenceDiagram
    participant W1 as Wafer 1
    participant W2 as Wafer 2
    participant W3 as Wafer 3
    W1->>W2: activaciones de la petición A
    W2->>W3: activaciones de la petición A
    W1->>W2: activaciones de la petición B
    W2->>W3: activaciones de la petición B
    Note over W1,W3: cada wafer procesa una petición distinta al mismo tiempo

💭 Clave: un wafer no espera a que termine la petición completa: procesa una porción distinta de cada petición al mismo tiempo que sus vecinos, así el pipeline nunca corre a lote 1.

Tensor: partir una misma capa en pedazos

El paralelismo de tensor divide una sola matriz de pesos en fragmentos, y cada chip calcula su fragmento de la multiplicación. Al final hay que combinar los resultados parciales con una operación de sincronización llamada all-reduce, lo que exige una red muy rápida entre chips (normalmente NVLink). Este es el enfoque que usa Megatron-LM para entrenar modelos tipo GPT en clústeres de GPU.

Expertos: solo activar una fracción del modelo

En modelos Mixture of Experts, cada token se enruta a un subconjunto de redes especializadas (expertos), y cada chip puede guardar solo algunos expertos en vez de una copia completa del modelo. El tráfico de red depende de a qué expertos se enrute cada token, así que no es predecible como en un pipeline fijo.

El caso Kimi K3 en los chips Cerebras WSE-3T

Kimi K3 tiene 69 capas de atención lineal y 24 capas densas, 93 en total. Si los 3 billones de parámetros se guardan en 4 bits, el peso crudo es de 1,5 TB; dividido entre 93 capas, cada una pesa alrededor de 16 GB, bien por debajo de los 44 GB que trae cada wafer.

Ahí aparece la diferencia entre tipos de capa. Una capa densa necesita espacio extra para el cache KV, porque guarda las claves y valores de cada token de cada petición activa; por eso entra solo una capa densa por wafer, dejando el resto de la memoria para ese cache. Una capa de atención lineal, en cambio, guarda un estado de tamaño fijo por petición sin importar cuánto contexto tenga, así que no necesita cache KV y caben dos por wafer (2 × 16 GB = 32 GB de 44 GB disponibles).

Con 69 capas lineales a dos por wafer hacen falta 35 wafers (34,5 redondeado hacia arriba), y con 24 capas densas a una por wafer hacen falta 24 wafers más. El total es 59 wafers WSE-3T. Como cada capa entera cabe en un solo wafer, los pesos nunca se dividen entre chips: solo las activaciones cruzan de wafer en wafer.

Según el mismo post de Cerebras, Kimi K3 anota 88,3% en Terminal-Bench 2.1 (contra 88,8% de GPT-5.6 Sol) y 91,2% en BrowseComp (contra 90,4% de GPT-5.6 Sol), resultados que lo ubican a la par de modelos propietarios en tareas de codificación y uso de herramientas, pese a ser un modelo de pesos abiertos.

Cómo empezar: probar el sharding de pesos con Hugging Face Accelerate

Nadie tiene wafers Cerebras en su escritorio, pero el mismo principio de reparto de capas se puede probar con herramientas abiertas. Antes de correr nada, instalá las dependencias (el comando es idéntico en Windows, macOS o Linux):

pip install transformers accelerate torch

Primero, cargá un modelo chico sin indicarle dónde vivir:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-1.5B-Instruct")
print(model.device)

Sin device_map ni una llamada a .to(), Transformers carga los pesos en memoria RAM del proceso. La salida es:

cpu

Ahora pedile a Accelerate que decida el reparto por vos:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-1.5B-Instruct",
    device_map="auto",
)
print(model.hf_device_map)

Con un modelo chico que entra entero en una sola GPU, Accelerate no reparte nada y la salida es simplemente:

{'': 0}

Ese diccionario vacío significa «todo el modelo en el dispositivo 0». Con un modelo que no entra en una sola GPU, el mismo atributo hf_device_map devuelve un mapeo capa por capa, algo del estilo {'model.embed_tokens': 0, 'model.layers.0': 0, 'model.layers.14': 1, 'lm_head': 1}: ahí sí se ve el sharding de pesos en acción, cada bloque de capas asignado a una GPU distinta.

💡 Tip: imprimí siempre model.hf_device_map después de cargar con device_map="auto": es la única forma directa de confirmar si Accelerate repartió el modelo o lo dejó entero en un dispositivo.

Casos de uso reales

• Cerebras + Kimi K3: inferencia en producción con pipeline de 59 wafers, cada uno con un grupo fijo de capas y memoria dedicada al cache KV.
• Entrenamiento con Megatron-LM: NVIDIA combina paralelismo de tensor, pipeline y datos en lo que se conoce como paralelismo 3D, para entrenar modelos tipo GPT en clústeres de miles de GPUs.
• DeepSpeed ZeRO: en vez de repartir capas, ZeRO-3 reparte los estados del optimizador y los gradientes entre GPUs, otra forma de resolver el mismo problema de memoria durante el entrenamiento.
• Servir con vLLM: el motor vLLM expone el flag --tensor-parallel-size para que un modelo grande se divida entre varias GPUs al momento de servir peticiones, sin tocar el código del modelo.

Errores comunes y buenas prácticas

• Confundir paralelismo de modelos con paralelismo de datos. El segundo copia el modelo completo en cada chip para atender más peticiones en paralelo; el primero reparte piezas de un único modelo demasiado grande para un chip. Son complementarios, no alternativas.
• No medir el tamaño real por capa. Una capa densa con cache KV pesa distinto que una capa de atención lineal sin cache; planear el reparto sin esa distinción sobreestima cuántos chips caben por wafer.
• Pipeline mal alimentado. Si solo hay una petición en vuelo, cada chip queda ocioso esperando su turno (el llamado pipeline bubble). La solución es mantener varias peticiones o microbatches corriendo a la vez.
• Ignorar la latencia de red entre chips. 2 microsegundos por salto suena insignificante, pero sumado a muchos saltos en un pipeline largo sí pesa en la latencia total de una petición.
• Asumir que device_map=»auto» siempre reparte algo. Con un modelo chico no hace sharding de pesos real: solo entra en acción cuando el modelo no cabe en un dispositivo.

Comparativa con alternativas

Ninguna técnica de paralelismo por pipeline, de tensor o de expertos gana siempre; cada una resuelve un cuello de botella distinto del hardware disponible.

Profundizando

El número de peticiones concurrentes que soporta un wafer con capa densa depende del cache KV disponible. Si el wafer tiene 44 GB y la capa densa usa 16 GB, quedan unos 28 GB libres para cache KV. La fórmula general es memoria de cache disponible dividida entre el tamaño del cache KV de una sola petición (que depende de la longitud de contexto, el número de cabezas de atención y la precisión del cache, 8 bits en el caso de Kimi K3). No hay forma de confirmar la cifra exacta de peticiones concurrentes de Kimi K3 sin ejecutar el modelo real: el extracto disponible del post de Cerebras describe el método pero no publica el número final.

La cantidad de peticiones en vuelo también depende de cómo se agrupan en microbatches. Según Cerebras, el número de peticiones en vuelo es aproximadamente la cantidad de etapas del pipeline multiplicada por el tamaño del microbatch: microbatches chicos bajan la latencia de cada petición individual, microbatches grandes suben el rendimiento total del sistema.

El siguiente diagrama contrasta cómo se estructura un paralelismo por pipeline frente a uno de tensor:

flowchart LR
    subgraph Pipeline["Paralelismo de pipeline"]
    P1["Capas 1 y 2 en el Chip A"] --> P2["Capas 3 y 4 en el Chip B"]
    end
    subgraph Tensor["Paralelismo de tensor"]
    T1["Mitad de la matriz en el Chip A"] --> T3["All-reduce final"]
    T2["Mitad de la matriz en el Chip B"] --> T3
    end

El paralelismo de tensor exige una red rapidísima entre chips, porque cada capa termina con una sincronización. El paralelismo por pipeline tolera redes más lentas, pero agrega latencia por petición, porque cada token tiene que atravesar todas las etapas antes de salir. La elección depende del cuello de botella real del hardware disponible, no hay una opción que gane siempre; por eso las tres variantes de paralelismo de modelos conviven en el mismo clúster cuando el modelo y el hardware lo justifican.

📖 Resumen en Telegram: Ver resumen

Tu próximo paso: instalá transformers y accelerate, cargá un modelo abierto con device_map=»auto» e imprimí model.hf_device_map para ver con tus propios ojos cómo Accelerate decide el reparto en tu hardware.

Preguntas frecuentes

¿El paralelismo de modelos reemplaza al paralelismo de datos?

No, son complementarios. El paralelismo de datos copia el modelo completo en cada chip para atender más peticiones en paralelo; el paralelismo de modelos reparte piezas de un único modelo demasiado grande para un chip. Sistemas grandes como Megatron-LM combinan ambos.

¿Por qué Kimi K3 necesita 59 wafers Cerebras WSE-3T?

Porque sus 93 capas no entran en un solo wafer de 44 GB. Cada capa pesa unos 16 GB en 4 bits, así que el reparto exige 35 wafers para las 69 capas de atención lineal y 24 wafers más para las capas densas.

¿device_map=»auto» de Hugging Face hace sharding de pesos real?

Sí, aplica el mismo principio a menor escala: reparte las capas del modelo entre las GPUs disponibles cuando el modelo no entra en una sola. Con un modelo chico, en cambio, no reparte nada.

¿El paralelismo por pipeline hace que un chip espere a que termine la petición anterior?

No. Cada etapa del pipeline procesa una petición distinta al mismo tiempo que sus vecinos, así un pipeline bien alimentado nunca corre a lote 1.

¿Qué diferencia hay entre paralelismo de tensor y paralelismo de expertos?

El de tensor divide una misma capa densa en pedazos que se ejecutan en paralelo con una sincronización final (all-reduce). El de expertos activa solo un subconjunto de redes especializadas por token, sin tocar el resto del modelo.

Referencias

• Cerebras: post oficial que detalla cómo reparten Kimi K3 en wafers WSE-3T.
• Hugging Face Accelerate: documentación de device_map=»auto» y el reparto automático de capas entre dispositivos.
• NVIDIA Megatron-LM: repositorio con la implementación de referencia del paralelismo de tensor y pipeline combinados.
• vLLM: documentación del motor de inferencia que expone –tensor-parallel-size para servir modelos grandes en producción.

📱 ¿Te gusta este contenido? Únete a nuestro canal de Telegram @programacion donde publicamos a diario lo más relevante de tecnología, IA y desarrollo. Resúmenes rápidos, contenido fresco todos los días.

Imagen destacada: Foto de Laura Ockel en Unsplash