El costo de verdad

Mixture of Experts, episodio diez, el último: el costo de verdad. Un MoE, mixture of experts, promete muchos parámetros con poco cómputo por token. Hoy vemos la letra chica: qué hay que tener en memoria, por qué el tamaño del batch cambia todo, cuánto viaja entre máquinas, y cómo se arranca un MoE sin entrenarlo desde cero.

Los cuatro costos

• Memoria: todos los expertos cargados, aunque se usen 8
• Batch: con pocos tokens, se lee mucho por token
• Comunicación: los tokens viajan a la placa de su experto
• Entrenamiento: upcycling, arrancar desde un denso

Cuatro costos. La memoria: aunque cada token use ocho expertos, los treinta y dos tienen que estar cargados, porque el próximo token puede elegir cualquiera. El batch: con pocos tokens a la vez, cada uno obliga a leer sus expertos, y leer pesos es lo que más tarda. La comunicación: cuando los expertos no entran en una sola placa, los tokens tienen que viajar a donde vive su experto. Y el entrenamiento: en lugar de empezar de cero, se puede reciclar un modelo denso. Vamos con números de Granite, el MoE abierto del episodio anterior.

Parte 1 · Memoria: se carga todo

Parámetros totales y parámetros activos son dos números distintos, y cada uno paga una cuenta distinta.

Totales contra activos en Granite

Contamos todos los parámetros. Separamos los de los expertos, que en Granite se llaman experts en el nombre. Dividimos por veinticuatro capas por treinta y dos expertos: el tamaño de uno. Los activos por token son todo lo que no es experto, más ocho expertos por capa. Y pasamos a bytes en bfloat16, el formato de dieciséis bits: dos bytes por número. Mil trescientos treinta y cinco millones de parámetros, y el noventa y uno por ciento está en los expertos. Por token se usan cuatrocientos veintinueve millones. En memoria: dos coma sesenta y siete gigabytes cargados, para usar cero coma ochenta y seis por token.

total = sum(p.numel() for p in model.parameters())
expertos = sum(p.numel() for n, p in model.named_parameters()
               if ".experts." in n)
por_experto = expertos / (24 * 32)            # 24 capas × 32 expertos
activos = total - expertos + 24 * 8 * por_experto   # 8 por token

print(f"totales : {total / 1e6:6.0f} M")
print(f"expertos: {expertos / 1e6:6.0f} M ({expertos / total:.0%})")
print(f"activos : {activos / 1e6:6.0f} M por token")
print(f"bf16    : {2 * total / 1e9:.2f} GB, {2 * activos / 1e9:.2f} por token")

totales :   1335 M
expertos:   1208 M (91%)
activos :    429 M por token
bf16    : 2.67 GB, 0.86 por token

La cuenta, término por término

gate y up: 2h × d; down: d × h

L = 24 capas, E = 32 expertos

k = 8: solo cambia E por k

La misma cuenta, en fórmula. Un experto de Granite es un MLP, multilayer perceptron, con compuerta: una matriz que proyecta hacia arriba y otra de compuerta, juntas de dos h por d, y una que baja, de d por h. d es mil veinticuatro y h, quinientos doce. En el dibujo, la proyección de entrada: mil veinticuatro filas por mil veinticuatro columnas. En total, tres por d por h: un millón quinientos setenta mil pesos. Los totales son el resto del modelo, atención y embeddings, más L capas por E expertos por ese tamaño. Los activos son la misma cuenta cambiando E por k, ocho. En el dibujo, las barras reales: la memoria paga los totales; el cómputo, los activos.

DeepSeek-AI, 2024 — DeepSeek-V3 Technical Report · arXiv:2412.19437

671B · 37B activos · «671B total parameters with 37B activated for each token».

A escala grande, la proporción es más extrema. DeepSeek-V3 tiene seiscientos setenta y un mil millones de parámetros, y por token activa treinta y siete mil millones: menos del seis por ciento. Y para servirlo hay que tenerlos todos cargados.

Parte 2 · El batch lo cambia todo

Cuando generás texto de a un token, la mayor parte del tiempo se va en leer pesos de la memoria, no en multiplicar.

¿Cuántos expertos distintos toca un batch?

al azar y parejo

n tokens independientes: se multiplica

E expertos, cada uno tocado con esa probabilidad

Si el router eligiera al azar y parejo, ¿cuántos expertos distintos toca un batch de n tokens? Un token no elige a un experto dado con probabilidad uno menos k sobre E: veinticuatro de treinta y dos, tres cuartos. Si los n tokens eligen de forma independiente, las probabilidades se multiplican: tres cuartos a la n. En el dibujo, cae rapidísimo: con dieciséis tokens, uno por ciento. El valor esperado de expertos distintos, esa E de doble trazo que significa promedio sobre todos los sorteos posibles, es E por uno menos eso. En el dibujo: un token toca ocho; cuatro tokens, casi veintidós; dieciséis, casi todos. Los pesos de cada experto tocado se leen una vez por batch y se reparten entre sus tokens.

Expertos tocados y megabytes por token

Ahora con el ruteo real de la capa doce de Granite. Para cada tamaño de batch, sorteamos doscientos batches de n tokens y contamos los expertos distintos. Al lado, la fórmula del azar. Y los megabytes de pesos que hay que leer por token, en tres casos: el MoE; un denso con el tamaño de los ocho expertos activos; y un denso con el tamaño de los treinta y dos. Con un token, el MoE lee veinticinco megabytes por token en esta capa, igual que el denso chico, y cuatro veces menos que el grande. Con dieciséis tokens toca veintinueve expertos, un poco menos que el azar. Y desde sesenta y cuatro, los toca todos: lee lo mismo que el denso grande, uno coma cinco siete, y con doscientos cincuenta y seis, cero coma treinta y nueve.

E, k = 32, 8
mb_exp = 3 * 512 * 1024 * 2 / 1e6            # MB de un experto en bf16
print("  n  dist  azar    MoE  d=act  d=tot")   # MB por token y capa
for n in [1, 4, 16, 64, 256]:
    distintos = []
    for _ in range(200):                       # 200 batches de n tokens
        fila = torch.randperm(len(idx), generator=g)[:n]
        distintos.append(idx[fila].unique().numel())
    d = sum(distintos) / len(distintos)
    azar = E * (1 - (1 - k / E) ** n)
    moe, act, tot = d * mb_exp / n, k * mb_exp / n, E * mb_exp / n
    print(f"{n:3d} {d:5.1f} {azar:5.1f} {moe:6.2f} {act:6.2f} {tot:6.2f}")

  n  dist  azar    MoE  d=act  d=tot
  1   8.0   8.0  25.17  25.17 100.66
  4  19.9  21.9  15.63   6.29  25.17
 16  29.1  31.7   5.73   1.57   6.29
 64  31.9  32.0   1.57   0.39   1.57
256  32.0  32.0   0.39   0.10   0.39

MB de expertos leídos por token y capa

La lección, en tres filas. Con un token por vez, el MoE lee como un denso del tamaño de sus activos: barato. Con dieciséis, ya está a mitad de camino. Y con un batch grande, lee como un denso del tamaño total. O sea: con batch chico, un MoE se comporta como un modelo chico, aunque igual hay que tener cargados los totales. Con batch grande, el ahorro en lectura de memoria desaparece, y lo que queda es el ahorro en cuentas: cada token sigue multiplicando solo por ocho expertos.

Parte 3 · Expert parallelism y all-to-all

Cuando los expertos no entran en una placa, se reparten: cada placa guarda algunos, y los tokens viajan.

Una capa MoE repartida en placas

El esquema se llama expert parallelism. Cada placa de video, la GPU, graphics processing unit, tiene sus propios tokens y una parte de los expertos. Cada placa corre el router para sus tokens. Después viene un all-to-all: todas las placas les mandan a todas las demás los vectores de los tokens que eligieron a sus expertos. Cada placa corre sus expertos sobre lo que recibió. Un segundo all-to-all devuelve los resultados a la placa de origen. Y ahí se hace la suma pesada por g y se vuelve al residual stream, la corriente principal del modelo. Dos viajes por capa MoE, ida y vuelta.

tokens (placa 0) → router local → all-to-all: ida → exp. 0–7 (placa 0) → exp. 8–15 (placa 1) → exp. 16–23 (placa 2) → exp. 24–31 (placa 3) → all-to-all: vuelta → suma ponderada → residual stream

Cuánto viaja, con el ruteo real

Simulémoslo con el ruteo de la capa doce. Cuatro placas, ocho expertos cada una. Los mil cincuenta y cinco tokens se reparten en partes iguales: es data parallelism, cada placa con su pedazo. El experto e vive en la placa e dividido ocho. La matriz de tráfico cuenta cuántas asignaciones salen de cada placa hacia cada otra. La diagonal es lo que no viaja. Y pasamos a megabytes. Solo el veinte por ciento se queda en casa; el ochenta viaja. Trece coma ocho megabytes de ida y otro tanto de vuelta, en una sola capa.

D = 4                                          # 4 placas, 8 expertos cada una
T = len(idx)
origen = torch.arange(T) * D // T              # data parallel: tokens repartidos
destino = idx // 8                             # experto e vive en la placa e // 8
trafico = torch.zeros(D, D, dtype=torch.long)
for o in range(D):
    a_donde = destino[origen == o].reshape(-1)
    trafico[o] = torch.bincount(a_donde, minlength=D)
print("de \\ a ", " ".join(f"{d:5d}" for d in range(D)))
for o in range(D):
    print(f"placa {o}", " ".join(f"{x:5d}" for x in trafico[o].tolist()))
local = trafico.diag().sum() / trafico.sum()
print(f"se queda: {local:.0%} · viaja: {1 - local:.0%}")
mb = T * 8 * 1024 * 2 / 1e6                    # T·k vectores de d=1024 en bf16
print(f"ida {mb * (1 - local):.1f} de {mb:.1f} MB, igual de vuelta")

de \ a      0     1     2     3
placa 0   409   437   609   657
placa 1   397   417   639   659
placa 2   484   669   546   413
placa 3   473   748   561   322
se queda: 20% · viaja: 80%
ida 13.8 de 17.3 MB, igual de vuelta

Matriz de tráfico, capa 12

La matriz. Cada fila es una placa de origen y cada columna, una de destino. Las celdas fuera de la diagonal son tráfico por la red. Fijate además que las columnas no suman lo mismo: la placa uno recibe dos mil doscientas setenta y una asignaciones y la cero, mil setecientas sesenta y tres. La placa más cargada marca el ritmo, porque todas esperan a la más lenta antes del all-to-all de vuelta. Por eso el balanceo de carga importa: es tiempo de reloj.

Cuánto viaja, en fórmula

cada token va a k expertos

un vector de d = 1024 en bf16 por asignación

parejo: con D = 4, el 75 %; medido: 80 %

La cuenta del tráfico. Cada token genera k asignaciones: mil cincuenta y cinco por ocho, ocho mil cuatrocientas cuarenta. Cada asignación es un vector de d números, dos bytes cada uno: diecisiete coma tres megabytes. En el dibujo, cómo se parte: tres y medio se quedan, trece coma ocho viajan. Si el ruteo fuera parejo, lo que viaja sería uno menos uno sobre D: con cuatro placas, el setenta y cinco por ciento; con más placas, casi todo. Medimos ochenta. Por eso DeepSeek-V3 limita cada token a, como mucho, cuatro nodos, y así acota el tráfico entre máquinas.

Fedus, Zoph y Shazeer, 2021 — Switch Transformers · arXiv:2101.03961

all-to-all · «Because each core has its own expert, we do an all-to-all communication» para repartir los tokens entre expertos.

Este esquema está descrito en el paper de Switch Transformer. Cada núcleo del chip tiene su propio experto, y para llevar cada token al suyo se hace una comunicación all-to-all, más otra igual para traer los resultados.

Parte 4 · Upcycling: reciclar un denso

Entrenar un MoE desde cero es caro. ¿Y si arrancamos desde un modelo denso que ya existe?

Copiar el MLP en cada experto

La receta se llama sparse upcycling, y es corta. Cada experto arranca como una copia exacta del MLP del modelo denso. El router es nuevo, con pesos chicos al azar: una normal de media cero y desvío cero coma cero dos, como en el paper. Al arrancar, el MoE calcula exactamente lo mismo que el denso: todos los expertos son iguales, así que el MLP sale como factor común, y los pesos g de los dos elegidos suman uno. Y contamos parámetros: totales y activos. La diferencia máxima es uno coma cinco por diez a la menos siete: ruido de punto flotante. El MoE tiene ocho veces los parámetros del denso, y por token activa dos expertos, más el router.

Cada línea, su término:

• experto.load_state_dict(denso.state_dict()) → $E_i \leftarrow \mathrm{MLP}_{\text{denso}} \quad \forall i$ — 8 copias idénticas
• init.normal_(moe.router.weight, std=0.02) → $W_r \sim \mathcal{N}(0,\, 0{,}02^2)$ — router nuevo, al azar
• dif = (moe(x) - denso(x)) → $\textstyle\sum_{S} g_i\, E_i(x) = \mathrm{MLP}(x)$ — factor común MLP(x); los g suman 1
• activos por token → $N_{\text{tot}} = 8\,N,\quad N_{\text{act}} \approx 2\,N$ — E = 8, k = 2

moe = MoE(d, h, E=8, k=2)                      # 8 expertos, 2 por token
for experto in moe.expertos:
    experto.load_state_dict(denso.state_dict())   # copia exacta del denso
torch.nn.init.normal_(moe.router.weight, std=0.02)  # router nuevo, al azar

with torch.no_grad():
    dif = (moe(x) - denso(x)).abs().max().item()
print(f"máx |MoE(x) − denso(x)| = {dif:.1e}")
print("parámetros denso:", parametros(denso))
print("parámetros MoE  :", parametros(moe), "· activos por token:",
      2 * parametros(denso) + moe.router.weight.numel())

máx |MoE(x) − denso(x)| = 1.5e-07
parámetros denso: 2128
parámetros MoE  : 17152 · activos por token: 4384

Seguir entrenando: denso contra upcycled

En la tarea sintética de la serie, ocho grupos de datos con funciones distintas. El denso de partida ya entrenó ochocientos pasos. Opción uno: seguir entrenándolo ochocientos pasos más. Opción dos: hacerle upcycling. Arranca con exactamente la misma pérdida. Y lo entrenamos los mismos ochocientos pasos, con la pérdida de balanceo del episodio cuatro. El denso baja de cero coma setecientos cuarenta y tres a cero coma seiscientos sesenta y cuatro. El upcycled, a cero coma ciento sesenta y cuatro. Ojo: el MoE tiene ocho veces los parámetros y el doble de cómputo activo. Lo que muestra es que la copia no estorba: los expertos se separan solos.

print(f"denso de partida      : {mse(base):.3f}")
densa = copy.deepcopy(base)                   # opción 1: seguir con el denso
entrenar(densa, X, Y, pasos=800, semilla=1)

up = copy.deepcopy(base)                      # opción 2: upcycling
moe = MoE(d, h, E=8, k=2)
for e in moe.expertos:
    e.load_state_dict(base.capa.state_dict())
torch.nn.init.normal_(moe.router.weight, std=0.02)
up.capa = moe
print(f"upcycled, paso 0      : {mse(up):.3f}")   # = denso
entrenar(up, X, Y, pasos=800, semilla=1, alfa=0.01)
print(f"+800 pasos, denso     : {mse(densa):.3f}")
print(f"+800 pasos, upcycled  : {mse(up):.3f}")

denso de partida      : 0.743
upcycled, paso 0      : 0.743
+800 pasos, denso     : 0.664
+800 pasos, upcycled  : 0.164

Komatsuzaki et al., 2022 — Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints · arXiv:2212.05055

~50 % · Los modelos upcycled superan a sus versiones densas «using only ~50% of the initial dense pretraining sunk cost».

El paper. Con T5 y con Vision Transformers, los modelos reciclados superan a sus versiones densas gastando, en el entrenamiento extra, más o menos la mitad de lo que había costado preentrenar el denso. Y le ganan a MoE entrenados desde cero con el cien por ciento de ese presupuesto. Cada experto arranca como copia del MLP, y el router, al azar.

Cierre · La serie, en diez líneas

Diez episodios. Volvamos a la promesa del principio: muchos parámetros, poco cómputo por token.

Los diez episodios

El recorrido completo. Un MoE parte el MLP en expertos y usa unos pocos por token: los parámetros totales pagan memoria y los activos, cómputo. El router es una capa lineal con un softmax y un top-k. Una capa MoE se implementa agrupando tokens por experto y devolviéndolos a su lugar. Sin presión, el router colapsa: la pérdida de balanceo y la capacidad lo contienen. Switch usa un experto y la z-loss achica los logits. Los expertos finos multiplican las combinaciones y los compartidos se llevan lo común. El balanceo también se puede hacer con un sesgo, sin pérdida auxiliar. Por dentro, el router separa tipos de token, no temas. Y hoy, el costo de verdad.

Resumiendo

• Memoria: los totales (Granite, 2,67 GB); cómputo: los activos
• Batch chico: lee como un denso chico; grande, como uno total
• Expert parallelism: dos all-to-all por capa MoE
• Con 4 placas, el 80 % de las asignaciones viaja
• Upcycling: copiar el MLP en cada experto, router al azar
• Fin de la serie: todos los snippets corren en el lab

Papers del episodio

• DeepSeek-AI, 2024 — DeepSeek-V3 Technical Report
• Fedus, Zoph y Shazeer, 2021 — Switch Transformers
• Komatsuzaki et al., 2022 — Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints