Muchos parámetros, poco cómputo

Mixture of Experts. Episodio uno: muchos parámetros, poco cómputo. Un MoE, mixture of experts, es un modelo que tiene muchísimos parámetros pero que, para cada token, usa solo una parte. Hoy vemos de dónde sale esa idea, cuánto ahorra, y lo contamos en un modelo abierto de verdad.

La idea en cuatro frases

• Un modelo denso usa todos sus pesos para cada token
• Un MoE parte una pieza en E expertos
• Un router elige k de esos E para cada token
• Parámetros totales ≠ parámetros activos

En un modelo denso, cada token pasa por todos los pesos: si tiene mil millones de parámetros, cada token hace cuentas con los mil millones. Un MoE toma una pieza del modelo y la parte en E copias independientes, que se llaman expertos. Un router, una capa chiquita, mira cada token y elige k de esos E expertos. Solo esos k hacen cuentas. Resultado: el modelo tiene muchos parámetros totales, pero cada token usa pocos parámetros activos. A esto se le dice cómputo condicional: qué se calcula depende de la entrada.

El mapa de la serie

El recorrido de la serie. Hoy, la idea y la cuenta. Después el mecanismo: cómo decide el router, y una capa escrita desde cero en PyTorch. Luego, cómo se rompe: el router que manda todo a pocos expertos, los tokens que no entran, y los trucos de estabilidad. Las variantes modernas: expertos chicos, expertos compartidos, y balanceo sin pérdida extra. Y al final abrimos un modelo real para ver qué aprendió cada experto, y cuánto cuesta servirlo.

Pocos activos → El router → Capa desde cero → Colapso → Capacidad → Switch · z-loss → Expertos finos → Sin pérdida aux → Qué aprenden → Costo real

Parte 1 · Dónde están los parámetros

Antes de partir algo, hay que saber qué conviene partir. En un transformer, casi todo el peso está en un solo lugar.

Un bloque denso, contado

Un bloque de transformer tiene dos piezas: la atención, donde los tokens se miran entre sí, y el MLP, multilayer perceptron, una red densa que procesa cada token por separado. En papers también se lo llama FFN, feed-forward network. d es el ancho del vector de cada token: mil veinticuatro. La atención de PyTorch, sin sesgos. El MLP sube de d a cuatro d, aplica GELU, una no linealidad suave, y baja de vuelta a d. Contamos con numel, que da la cantidad de números de cada tensor. La atención tiene cuatro millones; el MLP, ocho. Dos tercios del bloque son el MLP.

d = 1024                                     # ancho del residual stream
atencion = nn.MultiheadAttention(d, num_heads=16, bias=False)
mlp = nn.Sequential(nn.Linear(d, 4 * d, bias=False),   # sube a 4d
                    nn.GELU(),
                    nn.Linear(4 * d, d, bias=False))   # baja a d
n_at = sum(p.numel() for p in atencion.parameters())
n_mlp = sum(p.numel() for p in mlp.parameters())
total = n_at + n_mlp
print(f"atención: {n_at:>10,}  = 4·d²  ({n_at / total:.0%})")
print(f"MLP:      {n_mlp:>10,}  = 8·d²  ({n_mlp / total:.0%})")

atención:  4,194,304  = 4·d²  (33%)
MLP:       8,388,608  = 8·d²  (67%)

Por qué el MLP pesa el doble

atención: cuatro matrices cuadradas

sube a 4d y vuelve a d

el MLP: dos matrices de 4d × d

La misma cuenta en símbolos. d es un escalar, el ancho del vector de cada token. La atención tiene cuatro matrices: W Q, W K, W V y W O, las que arman queries, keys, values y la salida. Cada una es de d por d. En el dibujo, cuatro bloques de d al cuadrado: cuatro d cuadrado. El MLP multiplica x, un vector de d números, por W up, una matriz de cuatro d por d: en el dibujo, cada fila de W up por x da un número, y salen cuatro d números. Después W down baja de cuatro d a d. Dos matrices de cuatro d por d: ocho d cuadrado. El doble que la atención. Por eso, si queremos más parámetros, el MLP es el lugar natural para agregarlos.

Parte 2 · Partir el MLP en expertos

La idea de un MoE: en vez de un MLP grande, muchos MLP chicos, y usar solo algunos por token.

Denso contra MoE

Las dos versiones, lado a lado. La capa densa tiene un MLP, todos los tokens pasan por él, y cada parámetro que agregás es una cuenta más por token. La capa MoE tiene E MLP independientes, cada uno con sus propios pesos: esos son los expertos. Un router decide, para cada token, cuáles k usar. Si agregás expertos y dejás k fijo, el modelo tiene más parámetros, pero cada token hace las mismas cuentas.

La salida de una capa MoE

$g_i(x) = 0$ para los expertos no elegidos

lo que hay que guardar

lo que trabaja por token

x es el vector de un token. E sub i de x es la salida del experto i, otro vector del mismo tamaño. g sub i de x es un número, el peso que el router le da al experto i. La salida y es la suma de las salidas de los expertos, cada una multiplicada por su peso. El truco es que casi todos los g valen cero. En el dibujo, el token cero del snippet que viene: solo los expertos tres y seis tienen peso, cero coma cincuenta y ocho y cero coma cuarenta y dos, y los otros seis ni se calculan. Para guardar el modelo necesitás el router y los E expertos: diecisiete mil ciento cincuenta y dos números. Para un token, solo el router y k expertos: cuatro mil cuatrocientos ochenta.

Un MoE de juguete

Esta es la capa MoE de juguete que vamos a usar en toda la serie; en el episodio tres la escribimos línea por línea. Ocho expertos, dos por token, vectores de treinta y dos números. Un experto es un MLP de treinta y dos a treinta y dos a treinta y dos: dos mil ciento doce parámetros. El router es una matriz de treinta y dos por ocho: doscientos cincuenta y seis. Pasamos cuatro tokens al azar. Y miramos qué eligió el router para cada uno. Diecisiete mil ciento cincuenta y dos totales, cuatro mil cuatrocientos ochenta activos: una cuarta parte. Cada token eligió su par de expertos, con pesos que suman uno.

moe = MoE(d=32, h=32, E=8, k=2)          # 8 expertos, 2 por token
p_exp = parametros(moe.expertos[0])      # un experto: 32 → 32 → 32
p_router = parametros(moe.router)        # W_r: 32 × 8
print("un experto:", p_exp, "| router:", p_router)
print("totales:   ", p_router + 8 * p_exp)
print("activos:   ", p_router + 2 * p_exp, "por token")
x = torch.randn(4, 32)                   # 4 tokens
y = moe(x)
for t in range(4):
    e = moe.ultimo["idx"][t].tolist()
    g = moe.ultimo["pesos"][t].tolist()
    print(f"t{t} → {e}  g {g[0]:.2f} {g[1]:.2f}")

un experto: 2112 | router: 256
totales:    17152
activos:    4480 por token
t0 → [3, 6]  g 0.58 0.42
t1 → [3, 4]  g 0.64 0.36
t2 → [7, 2]  g 0.61 0.39
t3 → [7, 6]  g 0.53 0.47

Qué eligió cada token (salida real)

La salida, en una tabla. El token cero va a los expertos tres y seis. El uno, al tres y al cuatro. El dos, al siete y al dos. El tres, al siete y al seis. El router todavía no aprendió nada: sus pesos son al azar, y por eso los pesos andan cerca de mitad y mitad. En el episodio dos vemos cómo sale cada número de esta tabla.

Parte 3 · Un MoE de verdad: Granite

Ahora la misma cuenta en un modelo abierto: Granite 3.1, de IBM, versión de mil millones de parámetros con cuatrocientos millones activos.

Contar sin cargar el modelo

No hace falta cargar el modelo para contar sus parámetros. Una función que agrupa cada tensor por el nombre que tiene en el archivo. Buscamos los archivos en la caché local, sin internet. safetensors es un formato que guarda, al principio del archivo, un índice con el nombre y la forma de cada tensor. Leemos solo esas formas y multiplicamos sus dimensiones. Granite usa ocho expertos de treinta y dos: los activos son el total menos las tres cuartas partes de los expertos. Mil doscientos ocho millones están en los expertos. Atención, setenta y cinco. Total: mil trescientos treinta y cinco millones. Activos por token: cuatrocientos veintinueve.

def grupo(nombre):
    if "router" in nombre: return "router"
    if "block_sparse_moe" in nombre: return "expertos"
    if "self_attn" in nombre: return "atención"
    return "embeddings" if "embed" in nombre else "normas"

carpeta = snapshot_download("ibm-granite/granite-3.1-1b-a400m-base",
                            local_files_only=True)
cuenta = Counter()
for archivo in glob(f"{carpeta}/*.safetensors"):
    with safe_open(archivo, "pt") as f:        # lee solo los headers
        for nombre in f.keys():
            forma = f.get_slice(nombre).get_shape()
            cuenta[grupo(nombre)] += math.prod(forma)
total = sum(cuenta.values())
activos = total - cuenta["expertos"] * (1 - 8 / 32)   # top-8 de 32
for g, n in cuenta.most_common():
    print(f"{g:<11}{n / 1e6:>8.1f} M")
print(f"total {total / 1e6:.1f} M | activos {activos / 1e6:.1f} M")

expertos     1208.0 M
atención       75.5 M
embeddings     50.3 M
router          0.8 M
normas          0.1 M
total 1334.6 M | activos 428.7 M

Activos en Granite

24 capas, 32 expertos por capa

$k = 8,\; E = 32$: se apagan 3/4 de los expertos

cada token usa un tercio del modelo

Desglosado. Cada experto de Granite tiene un millón quinientos setenta y dos mil parámetros. Por treinta y dos expertos y por veinticuatro capas: mil doscientos ocho millones. En el dibujo, la barra de los expertos se come casi todo el total. Para un token, de los expertos solo trabajan ocho de treinta y dos, una cuarta parte: trescientos dos millones. La atención, los embeddings, el router y las normas se usan siempre. Contamos los embeddings porque en Granite la misma matriz se usa también a la salida, para sacar los logits. En el dibujo, la suma da cuatrocientos veintiocho coma siete millones. Un token usa un tercio del modelo.

Un experto de Granite por dentro

¿Y qué es un experto? Un MLP, con una vuelta. Leemos solo el experto cero de la capa cero. W in, de mil veinticuatro por mil veinticuatro, produce dos vectores de quinientos doce de una sola vez: gate y up. gate pasa por SiLU, una no linealidad suave, y multiplica a up número por número: es una compuerta. W out baja de quinientos doce a mil veinticuatro. El token entra con mil veinticuatro números y sale con mil veinticuatro. Un experto tiene un millón quinientos setenta y dos mil parámetros; por treinta y dos y por veinticuatro capas, los mil doscientos ocho millones de antes.

capa = "model.layers.0.block_sparse_moe"
for archivo in glob.glob(f"{carpeta}/*.safetensors"):
    with safe_open(archivo, "pt") as f:
        if f"{capa}.input_linear.weight" in f.keys():
            W_in = f.get_slice(f"{capa}.input_linear.weight")[0]    # experto 0
            W_out = f.get_slice(f"{capa}.output_linear.weight")[0]
print("W_in", tuple(W_in.shape), "W_out", tuple(W_out.shape))
x = torch.randn(1024, dtype=W_in.dtype)        # un token
gate, up = (W_in @ x).chunk(2)                 # 1024 → 512 y 512
h = F.silu(gate) * up                          # compuerta: 512
y = W_out @ h                                  # 512 → 1024
print("x", tuple(x.shape), "→ h", tuple(h.shape), "→ y", tuple(y.shape))
n = W_in.numel() + W_out.numel()
print(f"un experto: {n:,}")
print(f"× 32 × 24 capas = {n * 32 * 24:,}")

W_in (1024, 1024) W_out (1024, 512)
x (1024,) → h (512,) → y (1024,)
un experto: 1,572,864
× 32 × 24 capas = 1,207,959,552

El experto de Granite

El mismo experto en cajas. El token, mil veinticuatro números, se multiplica por W in. La salida se parte en dos mitades de quinientos doce: una decide cuánto pasa, la otra es lo que pasa. Y W out vuelve al ancho del modelo. Es la misma forma que el MLP denso, solo que mucho más angosto: quinientos doce por dentro en vez de cuatro mil. Muchos expertos angostos en vez de uno ancho.

x: 1024 → W_in: 1024×1024 → gate 512 · up 512 → SiLU(gate) · up → W_out: 1024×512 → y: 1024

Parte 4 · Memoria contra cómputo

Pocos parámetros activos quiere decir pocas cuentas. Pero no quiere decir poca memoria.

Qué paga cada uno

cada peso: una multiplicación y una suma

cuentas: dependen de los activos

bfloat16: todos los expertos, siempre cargados

Separemos los dos costos. Cada peso que participa en un producto de matrices hace una multiplicación y una suma. En el dibujo, una neurona con tres entradas: tres productos, que se suman en z. Por eso las cuentas por token, que se miden en FLOPs, floating point operations, son más o menos dos veces los parámetros activos. Granite hace cero coma ochenta y seis miles de millones por token; un modelo denso del mismo tamaño haría dos coma sesenta y siete. Pero la memoria depende de los totales: no sabés qué experto va a pedir el próximo token, así que tienen que estar todos cargados. En bfloat16, un formato de dos bytes por número.

Las dos cuentas

En código es aritmética. Dos bytes por número. La memoria de los pesos usa el total, los mil trescientos treinta y cinco millones. Las cuentas por token usan los activos. Y comparamos con un denso hipotético del mismo tamaño total. Dos coma sesenta y siete gigabytes de pesos. Cero coma ochenta y seis contra dos coma sesenta y siete miles de millones de cuentas: el MoE hace el treinta y dos por ciento. Ojo, esto no cuenta la atención sobre el contexto, que crece con el largo del texto, ni el tiempo de mover datos. El costo de verdad lo medimos en el último episodio.

bytes_bf16 = 2                                 # un número en bfloat16
memoria = total * bytes_bf16                   # hay que tener TODO cargado
flops_moe = 2 * activos                        # por token: 1 mult + 1 suma
flops_denso = 2 * total                        # un denso del mismo tamaño
print(f"memoria de pesos:  {memoria / 1e9:.2f} GB")
print(f"FLOPs/token MoE:   {flops_moe / 1e9:.2f} G")
print(f"FLOPs/token denso: {flops_denso / 1e9:.2f} G")
print(f"el MoE hace {flops_moe / flops_denso:.0%} de las cuentas")

memoria de pesos:  2.67 GB
FLOPs/token MoE:   0.86 G
FLOPs/token denso: 2.67 G
el MoE hace 32% de las cuentas

Shazeer et al., 2017 — Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer · arXiv:1701.06538

137B · «greater than 1000x improvements in model capacity with only minor losses in computational efficiency»; una capa MoE de hasta 137.000 millones de parámetros entre capas LSTM.

La idea no es nueva, pero este paper la hizo funcionar a escala. Shazeer y su equipo, en Google, en dos mil diecisiete. Pusieron una capa MoE con hasta ciento treinta y siete mil millones de parámetros entre capas LSTM, long short-term memory, las redes recurrentes de esa época. Dicen haber logrado más de mil veces más capacidad con pérdidas menores de eficiencia de cómputo. La mayoría de los MoE de hoy usan la misma receta, cambiando el MLP de un transformer.

Lo que esta cuenta esconde

• Alguien tiene que elegir: el router
• Los expertos pueden quedar sin trabajo o con demasiado
• Cada experto recibe un pedazo del batch
• La memoria no baja: hay que tener todo cargado

La cuenta es linda, pero esconde los problemas que ocupan el resto de la serie. Alguien tiene que elegir qué expertos usa cada token, y esa elección se aprende: es el router. Si el router se enamora de dos expertos, los demás no aprenden nada y el modelo grande se vuelve uno chico. Cada experto recibe solo los tokens que le tocan, así que trabaja con un batch más chico y la placa rinde menos. Y la memoria sigue siendo la del modelo entero.

Resumiendo

• El MLP tiene 2/3 de un bloque: 8d² contra 4d²
• MoE: E expertos, un router elige k por token
• $y = \sum_i g_i(x)\,E_i(x)$, casi todos los $g_i$ en cero
• Granite: 1335 M totales, 429 M activos
• Cuentas ∝ activos; memoria ∝ totales
• Próximo: el router, número por número

Papers del episodio

• Shazeer et al., 2017 — Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer