Una capa MoE desde cero

Mixture of Experts. Episodio tres: una capa MoE, mixture of experts, desde cero. Ya sabemos qué hace el router. Hoy escribimos la capa entera en PyTorch: primero de la forma obvia, que es lenta, y después de la forma en que se hace de verdad. Al final la entrenamos contra una red densa.

Lo que tiene que hacer la capa

• Rutear: k expertos y k pesos por token
• Dispatch: juntar los tokens de cada experto
• Una llamada por experto, con todos sus tokens
• Combine: devolver cada salida a su token, pesada

El trabajo de la capa, en cuatro pasos. Rutear: el router le da a cada token sus k expertos y sus k pesos, como vimos en el episodio dos. Dispatch, que en inglés es despachar: juntar, para cada experto, todos los tokens que lo eligieron. Correr cada experto una sola vez, con todos sus tokens juntos. Y combine: devolver cada salida al token del que vino, multiplicada por su peso, y sumarla. Los nombres dispatch y combine son los que vas a encontrar en todo el código de MoE.

Parte 1 · La versión ingenua

La forma más directa de escribir la fórmula: un loop por token, y adentro un loop por experto elegido.

Un loop por token

Usamos la capa de juguete: ocho expertos, dos por token, vectores de treinta y dos. Primero el router, con la función del episodio anterior: pesos e índices, cada uno de forma T por k. Recorremos los tokens uno por uno. Para cada token, sus dos expertos. Llamamos al experto con ese único token y sumamos su salida, pesada. Es la fórmula de la suma, al pie de la letra. Y contamos cuántas veces llamamos a un experto. Mil veinticuatro tokens, dos mil cuarenta y ocho llamadas. Cada llamada multiplica una matriz por un solo vector. Este loop no se tira: queda como referencia. Cualquier versión rápida que escribamos tiene que dar exactamente lo mismo que él.

def moe_loop(x, moe):
    pesos, idx, _ = rutear(moe.router(x), moe.k)    # [T, k] y [T, k]
    y = torch.zeros_like(x)
    llamadas = 0
    for t in range(x.shape[0]):                      # cada token
        for j in range(moe.k):                       # cada elegido
            e = idx[t, j].item()
            y[t] += pesos[t, j] * moe.expertos[e](x[t])
            llamadas += 1
    return y, llamadas

with torch.no_grad():
    y, llamadas = moe_loop(x, moe)
print("tokens", x.shape[0], "| llamadas", llamadas)

tokens 1024 | llamadas 2048

Muchas llamadas chicas contra pocas grandes

¿Por qué es lento? El loop hace dos mil cuarenta y ocho llamadas, cada una con un solo vector. Una placa de video, una GPU, graphics processing unit, rinde cuando multiplica matrices grandes de una vez; con un vector por llamada pasa más tiempo esperando que calculando. Agrupando por experto hacemos solo ocho llamadas, una por experto, y cada una procesa todos sus tokens juntos: en promedio, T por k dividido E. Dos mil cuarenta y ocho dividido ocho: doscientos cincuenta y seis vectores por llamada.

Shazeer et al., 2017 — Outrageously Large Neural Networks, sección 3.1 · arXiv:1701.06538

kb/n ≪ b · «for a batch of b examples, each expert receives a much smaller batch of approximately kb/n ≪ b examples»: el shrinking batch problem.

Ese promedio tiene nombre. El paper de Shazeer lo llama shrinking batch problem, el problema del batch que se achica: si el router elige k de n expertos, cada experto recibe aproximadamente k por b dividido n ejemplos de un batch de b, mucho menos que b. Por eso agrupar importa tanto: es la única manera de que cada experto trabaje con un bloque grande de tokens.

Parte 2 · Dispatch: agrupar por experto

El truco para agrupar sin un loop por token es un ordenamiento. Lo vemos con seis tokens.

Dispatch y combine, en cajas

El plan completo. Partimos de idx, la tabla de T por k con los expertos de cada token, y la aplanamos: una fila por asignación, es decir, por cada par token y experto. Ordenamos esas filas por experto con argsort, y las cortamos en pedazos según cuántas tiene cada uno. Cada experto recibe su pedazo y hace una sola llamada. index add devuelve cada salida a la fila de su token, y queda y, del mismo tamaño que la entrada.

idx: T × k → aplanar: T·k filas → argsort por experto → split por cuenta → E llamadas → index_add_ → y: T × d

Ordenar para agrupar

Seis tokens, dos expertos por token, cuatro expertos; los índices son al azar. repeat interleave repite cada número de token dos veces: cero, cero, uno, uno, y así. Es de qué token viene cada fila. Aplanamos idx: a qué experto va cada fila. argsort da el orden que pone los expertos de menor a mayor; stable quiere decir que, dentro de un mismo experto, se respeta el orden de llegada. bincount cuenta cuántas filas hay de cada experto. Y split corta la lista de tokens en esos pedazos. Cuentas: dos, dos, tres y cinco. El experto dos recibe los tokens cero, uno y tres; el tres, cinco tokens.

T, k, E = 6, 2, 4
idx = torch.stack([torch.randperm(E)[:k] for _ in range(T)])  # [T, k]
print("idx por token:", idx.tolist())
tok = torch.arange(T).repeat_interleave(k)    # de qué token es cada fila
exp = idx.reshape(-1)                         # a qué experto va
orden = exp.argsort(stable=True)              # agrupar por experto
tok, exp = tok[orden], exp[orden]
cuenta = torch.bincount(exp, minlength=E)     # cuántos por experto
for e, filas in enumerate(tok.split(cuenta.tolist())):
    print(f"experto {e} ({cuenta[e]}): tokens {filas.tolist()}")

idx por token: [[2, 3], [3, 2], [0, 3], [3, 2], [1, 3], [1, 0]]
experto 0 (2): tokens [2, 5]
experto 1 (2): tokens [4, 5]
experto 2 (3): tokens [0, 1, 3]
experto 3 (5): tokens [0, 1, 2, 3, 4]

Quién recibe qué (salida real)

La salida, ordenada. El experto cero recibe los tokens dos y cinco. El uno, el cuatro y el cinco. El dos, el cero, el uno y el tres. Y el tres, cinco de los seis tokens. Fijate que ya acá, con datos al azar, la carga no es pareja: un experto tiene más del doble de trabajo que otros. Ese desbalance es el tema del próximo episodio.

Dispatch y combine, en fórmulas

elegir las filas de los tokens del experto $e$

una sola llamada por experto

combine: cada salida vuelve a su token

Lo mismo en símbolos. X es la matriz de entrada: una fila por token, cada fila de d números. T de t es el conjunto de expertos que eligió el token t. t sub e es la lista de tokens que eligieron al experto e. X sub e son esas filas de X, nada más. En el dibujo, el experto dos se lleva las filas cero, uno y tres. Elegir filas es un índice, no una cuenta. El experto procesa esas filas juntas: en el dibujo, tres tokens multiplicados de una vez por la primera matriz del experto. Y para cada token, sumamos lo que le devolvieron sus expertos, cada uno con su peso g. En el dibujo, el token cero recibe del experto dos y del tres.

Parte 3 · Combine: devolver en orden

Ya sabemos agrupar. Ahora la capa entera, y la prueba de que da lo mismo que el loop.

La capa agrupada

Esta es la versión que va adentro de la capa de juguete. Aplanamos, y argsort ordena las asignaciones por experto; los pesos viajan con ellas. Cortamos tokens y pesos en un pedazo por experto. Cada experto se llama una sola vez, con todas sus filas. index add suma cada fila de la salida, multiplicada por su peso, en la fila del token del que vino. Si un token tiene dos expertos, su fila recibe dos sumas. El error máximo contra el loop es uno coma ocho por diez a la menos siete: redondeo de punto flotante, porque las sumas se hacen en otro orden. Mismo resultado, con ocho llamadas en lugar de dos mil cuarenta y ocho. Esta es, casi línea por línea, la capa de juguete que usamos en toda la serie.

Cada línea, su término:

• orden = exp.argsort(stable=True) → $\pi = \mathrm{argsort}(e_1, \dots, e_{Tk})$ — dispatch: agrupar por experto
• tok.split(cuenta), w.split(cuenta) → $t_e,\; g_e$ — los tokens y pesos de cada experto
• salida = moe.expertos[e](x[t]) → $Y_e = E_e(X[t_e])$ — una llamada por experto
• y.index_add_(0, t, g[:, None] * salida) → $y_{t_e[j]} \mathrel{+}= g_{e,j}\, Y_e[j]$ — combine: sumar en la fila del token

def moe_agrupado(x, moe):
    pesos, idx, _ = rutear(moe.router(x), moe.k)
    tok = torch.arange(x.shape[0]).repeat_interleave(moe.k)
    exp, w = idx.reshape(-1), pesos.reshape(-1)
    orden = exp.argsort(stable=True)
    tok, exp, w = tok[orden], exp[orden], w[orden]       # dispatch
    cuenta = torch.bincount(exp, minlength=moe.E).tolist()
    y = torch.zeros_like(x)
    for e, t, g in zip(range(moe.E), tok.split(cuenta), w.split(cuenta)):
        salida = moe.expertos[e](x[t])                    # 1 llamada
        y.index_add_(0, t, g[:, None] * salida)           # combine
    return y

with torch.no_grad():
    y1, _ = moe_loop(x, moe)                              # la referencia
    y2 = moe_agrupado(x, moe)
print(f"error máximo: {(y1 - y2).abs().max().item():.1e}")

error máximo: 1.8e-07

Una prueba de cordura: k = E

Una forma de confiar en una implementación: llevarla a un caso que sabemos calcular a mano. Ocho expertos y k igual a ocho: todos activos. Ponemos la matriz del router en cero, así todos los logits valen lo mismo. Calculamos aparte el promedio de los ocho expertos. Y comparamos. Cada peso vale cero coma ciento veinticinco, un octavo, y la diferencia con el promedio es cero. Con todos los expertos y un router que no prefiere a nadie, un MoE es simplemente una red densa: el promedio de E MLP, multilayer perceptron, en paralelo.

moe = MoE(d=32, h=32, E=8, k=8)                  # k = E: todos activos
with torch.no_grad():
    moe.router.weight.zero_()                    # logits iguales
    x = torch.randn(5, 32)
    y = moe(x)
    promedio = sum(e(x) for e in moe.expertos) / 8
print("pesos de un token:", moe.ultimo["pesos"][0].tolist()[:4], "...")
print(f"diferencia con el promedio: {(y - promedio).abs().max().item():.1e}")

pesos de un token: [0.125, 0.125, 0.125, 0.125] ...
diferencia con el promedio: 0.0e+00

Por qué da el promedio

softmax de números iguales: reparto parejo

renormalizar no cambia nada: ya suman 1

el promedio de los expertos

La cuenta, sin saltear pasos. Si todos los logits son iguales, cada exponencial es la misma, e a la ele. El softmax es esa exponencial dividida por E veces ella misma: uno sobre E. En el dibujo, ocho barras iguales de doce coma cinco por ciento. Con k igual a E elegimos a todos, y renormalizar divide por la suma de los E pesos, que ya era uno: no cambia nada. Así que la salida es la suma de los E expertos, cada uno por un octavo. En el dibujo, ocho columnas con el mismo peso: un promedio.

Parte 4 · ¿Sirve?

La capa funciona. La pregunta que importa es otra: con las mismas cuentas por token, ¿un MoE aprende más que una red densa?

El experimento

• Datos sintéticos: 8 regiones, cada una con su función
• MoE: 8 expertos de 32, top-2
• Denso 64: mismos parámetros activos
• Denso 256: mismos parámetros totales

El experimento, con datos inventados. Cada ejemplo cae en una de ocho regiones del espacio, y cada región tiene su propia función no lineal para predecir. El MoE: ocho expertos de treinta y dos neuronas, dos por token. Un denso con sesenta y cuatro neuronas: hace las mismas cuentas por token que el MoE, porque dos expertos de treinta y dos son sesenta y cuatro. Y un denso de doscientas cincuenta y seis neuronas, con casi los mismos parámetros totales que el MoE. Los tres se entrenan mil quinientos pasos, con tres semillas, y medimos el error en ejemplos que nunca vieron.

Los resultados

El entrenamiento tarda unos segundos por modelo, así que lo corrimos una vez en un script aparte y acá leemos los resultados. El archivo que guardó ese script. Promediamos el error de test de las tres semillas. Es el MSE, mean squared error: el promedio de los errores al cuadrado. Y mostramos parámetros totales y activos de cada capa. El MoE: diecisiete mil totales, cuatro mil quinientos activos, error cero coma ciento cuarenta y ocho. El denso con los mismos activos: cero coma setecientos veinticuatro, casi cinco veces peor. Y el denso con los mismos totales: cero coma trescientos ochenta.

# entrenados una vez con lab/preparar_ep03.py (3 semillas, 1500 pasos)
r = json.load(open(os.path.join(DATOS, "ep03_resultados.json")))
print(f"{'modelo':<14}{'totales':>8}{'activos':>8}{'MSE test':>10}")
for nombre, v in r.items():
    mse = sum(v["test"]) / len(v["test"])
    print(f"{nombre:<14}{v['params_capa']:>8}{v['activos_capa']:>8}{mse:>10.3f}")

modelo         totales activos  MSE test
moe_8x32_top2    17152    4480     0.148
denso_64          4192    4192     0.724
denso_256        16672   16672     0.380

Error de test, promedio de 3 semillas

En barras, más bajo es mejor. El MoE queda abajo de todo. El denso que hace las mismas cuentas por token tiene casi cinco veces más error. Y lo llamativo: el denso con los mismos parámetros totales, que hace casi cuatro veces más cuentas, también pierde. En esta tarea, partir el problema entre expertos ayuda por sí mismo.

• MoE 8×32, top-2: 0,148
• denso 64 (mismos activos): 0,724
• denso 256 (mismos totales): 0,380

Ojo con esta comparación

• La tarea tiene regiones: favorece a los expertos a propósito
• Es un juguete: 1500 pasos, una capa
• En modelos reales se compara a igual cómputo de entrenamiento
• Lo que sí muestra: más capacidad sin más cuentas

Antes de sacar conclusiones. La tarea tiene ocho regiones con funciones distintas: está armada para que dividir el trabajo sirva. En datos sin esa estructura la ventaja puede ser chica o nula. Es un juguete: una sola capa, mil quinientos pasos. En modelos de lenguaje, las comparaciones serias se hacen a igual cómputo de entrenamiento, y las vemos con los papers en los próximos episodios. Lo que sí muestra es la mecánica: el MoE agregó capacidad sin agregar cuentas por token.

Resumiendo

• Loop por token: T·k llamadas de un vector, lento
• Dispatch: aplanar, argsort por experto, split
• Una llamada por experto con todos sus tokens
• Combine: index_add_ con el peso de cada asignación
• k = E y router en cero → promedio denso (verificado)
• Próximo: el router que colapsa y la pérdida de balanceo

Papers del episodio

• Shazeer et al., 2017 — Outrageously Large Neural Networks