Capacidad y tokens descartados

Mixture of Experts. Episodio cinco: capacidad y tokens descartados. En una capa MoE, mixture of experts, cada experto recibe una cantidad de tokens que no se sabe de antemano. El hardware, en cambio, quiere saberla. Hoy vemos la solución de compromiso: un cupo fijo por experto, qué pasa con los tokens que no entran, y cuánto cuesta.

Por qué un cupo fijo

• Los compiladores de TPU quieren formas fijas
• Entre placas, se manda un bloque de tamaño conocido
• La memoria se reserva antes de saber la carga
• Solución: C lugares por experto, ni uno más

En nuestra capa de juguete, cada experto recibía los tokens que le tocaran, fueran tres o mil. En hardware de verdad eso es incómodo. Las TPU, tensor processing unit, los chips de Google para redes neuronales, compilan el programa entero y necesitan que cada tensor tenga una forma conocida de antemano. Cuando los expertos viven en placas distintas, los tokens viajan entre ellas en bloques, y conviene que esos bloques tengan siempre el mismo tamaño. La memoria también se reserva antes de saber cómo cae la carga. La salida clásica: cada experto tiene C lugares. Ni uno más.

Parte 1 · La capacidad

¿Cuántos lugares por experto? El reparto perfecto, más un margen.

El capacity factor

lo que le tocaría a cada uno si el reparto fuera perfecto

$\lceil \cdot \rceil$: redondear para arriba

con factor > 1 sobran lugares: relleno

Con los números del episodio anterior: T, dos mil cuarenta y ocho tokens de test; k, dos; E, ocho expertos. T por k es la cantidad de asignaciones: cuatro mil noventa y seis. Divididas parejo entre ocho expertos, quinientas doce cada uno. En el dibujo, ocho barras iguales. La capacidad C es eso multiplicado por un número que se elige, el capacity factor, y redondeado para arriba: eso significan esos corchetes abiertos arriba. Con factor uno, quinientos doce; con uno coma veinticinco, seiscientos cuarenta; con dos, mil veinticuatro. En total se reservan E por C lugares. Si el factor es mayor que uno, sobran: los lugares vacíos se rellenan con ceros y se calculan igual.

Lugares reservados

La misma cuenta en código. Tokens, expertos por token y expertos. La capacidad, con ceil, que redondea para arriba. Y los lugares totales. Con factor uno, cuatro mil noventa y seis lugares para cuatro mil noventa y seis asignaciones: alcanza solo si el reparto es perfecto. Con dos, el doble de lugares: la mitad de las cuentas de los expertos se haría sobre relleno, aun con un router perfecto.

T, k, E = 2048, 2, 8                     # tokens del batch, top-k, expertos
for factor in (1.0, 1.25, 2.0):
    C = math.ceil(factor * T * k / E)    # capacidad de cada experto
    lugares = E * C                      # filas que se reservan
    print(f"factor {factor:<4}: C = {C:>4} | lugares {lugares:>5}"
          f" para {T * k} asignaciones")

factor 1.0 : C =  512 | lugares  4096 para 4096 asignaciones
factor 1.25: C =  640 | lugares  5120 para 4096 asignaciones
factor 2.0 : C = 1024 | lugares  8192 para 4096 asignaciones

Fedus, Zoph & Shazeer, 2021 — Switch Transformers, ecuación 3 · arXiv:2101.03961

factor > 1,0 · «A capacity factor greater than 1.0 creates additional buffer to accommodate for when tokens are not perfectly balanced across experts.»

Así lo define Switch Transformer. Su fórmula es tokens por batch dividido por expertos, por el capacity factor, sin la k, porque Switch manda cada token a un solo experto; con k igual a uno es la nuestra. Un factor mayor que uno crea un margen para cuando los tokens no se reparten perfecto. Y advierten el costo: valores altos desperdician cómputo y memoria.

Parte 2 · Overflow: el token que no entra

Si a un experto le llegan más tokens que lugares, los que llegan tarde quedan afuera. Lo vemos con ocho tokens.

Un buffer de forma fija

Ocho tokens, cuatro expertos, uno por token, para que se vea fácil. El experto que eligió cada token: el cero es el favorito, con cuatro tokens. Capacidad: ocho dividido cuatro, dos lugares por experto. Después de agrupar por experto, como en el episodio tres, calculamos el puesto de cada token en la fila de su experto: primero, segundo, tercero. Entra si su puesto es menor que C. Y el buffer tiene forma fija: E por C por d, pase lo que pase. El experto cero se queda con los tokens cero y dos; el tres y el seis llegaron tarde y se descartan. El uno y el tres tienen un lugar vacío, con menos uno.

T, k, E, d, factor = 8, 1, 4, 3, 1.0
x = torch.randn(T, d)
exp = torch.tensor([0, 2, 0, 0, 1, 2, 0, 3])      # experto de cada token
C = math.ceil(factor * T * k / E)                  # = 2
orden = exp.argsort(stable=True)                   # agrupar, por llegada
cuenta = torch.bincount(exp, minlength=E)
pos = torch.arange(T) - (cuenta.cumsum(0) - cuenta)[exp[orden]]
entra = pos < C                                    # cabe en su fila
buffer = torch.zeros(E, C, d)                      # forma fija
quien = torch.full((E, C), -1)                     # -1 = lugar vacío
e, p, t = exp[orden][entra], pos[entra], orden[entra]
buffer[e, p], quien[e, p] = x[t], t
print("buffer:", tuple(buffer.shape), "| token en cada lugar:")
print(quien.tolist())
print("descartados:", orden[~entra].tolist())

buffer: (4, 2, 3) | token en cada lugar:
[[0, 2], [4, -1], [1, 5], [7, -1]]
descartados: [3, 6]

El buffer, lugar por lugar (salida real)

El buffer, como tabla. El experto cero llena sus dos lugares y le sobran dos tokens. El uno tiene un lugar vacío, que igual se calcula. El dos llena justo. El tres, otro vacío. Fijate la paradoja: hay dos lugares libres y dos tokens afuera, pero los lugares son de otros expertos. La capacidad es por experto, no total.

Quién entra y qué le pasa al que no

por orden de llegada dentro de cada experto

la conexión residual de la red

el token sigue de largo, sin cambios

En fórmulas. Cada asignación tiene un puesto en la fila de su experto, según el orden de llegada. Entra si el puesto es menor que C; el corchete vale uno o cero. En el dibujo, las cuatro del experto cero: puestos cero y uno entran, dos y tres quedan afuera. ¿Y qué pasa con un token afuera? La capa MoE está metida en una conexión residual: la salida es h más lo que devuelva la capa. h es el vector del token, de d números. En el dibujo, la columna de la capa MoE está apagada. Si ningún experto lo procesó, la capa devuelve ceros, y h prima es igual a h: el token pasa de largo, como si esta capa no existiera para él.

Comprobarlo con un hook

Lo comprobamos en la red de juguete. Un capacity factor ridículo, cero coma veinticinco, para forzar descartes. Colgamos un hook: una función que PyTorch llama después del forward de un módulo, y que acá guarda la salida de la capa MoE. Contamos las filas de esa salida que son exactamente cero. Y la capacidad que resultó. Cuatro lugares por experto, treinta y dos en total, para ciento veintiocho asignaciones. Cuarenta y cinco de los sesenta y cuatro tokens no consiguieron ningún experto, y la capa devolvió cero exacto en esas cuarenta y cinco filas. Siguen solo por la residual.

red = Red(16, 32, 8, MoE(32, 32, E=8, k=2, factor_capacidad=0.25))
x = torch.randn(64, 16)
salidas = {}
red.capa.register_forward_hook(lambda m, i, o: salidas.update(moe=o))
with torch.no_grad():
    red(x)
u = red.capa.ultimo
sin = salidas["moe"].abs().sum(-1) == 0           # filas en cero
print("capacidad:", red.capa.capacidad(64), "por experto")
print("tokens sin experto:", u["tokens_sin_experto"], "de 64")
print("salida de la capa en cero para esos tokens:", int(sin.sum()))
# h = h + capa(norma(h)): esos tokens siguen con h intacto

capacidad: 4 por experto
tokens sin experto: 45 de 64
salida de la capa en cero para esos tokens: 45

Lepikhin et al., 2020 — GShard, sección 2.2 · arXiv:2006.16668

overflowed token · «When both experts selected by a token already exceed their capacity, the token is considered as an overflowed token […] passed on to the next layer via residual connections.»

Esta es la regla de GShard, de Google, el paper que llevó los MoE a más de seiscientos mil millones de parámetros. Usan top-2. Si los dos expertos de un token ya están llenos, el token se considera desbordado, overflowed, y su representación pasa a la capa siguiente por la conexión residual. Exactamente lo que acabamos de verificar.

Parte 3 · Cuánto se pierde

Ahora con los dos modelos entrenados del episodio cuatro: el que colapsó y el balanceado. Les ponemos un cupo solo al evaluar.

Descartes según el factor

Los dos modelos: sin balanceo y con alfa cero coma uno. Cuatro capacidades: sin límite, y capacity factor uno, uno coma veinticinco y dos. Cambiar el cupo es cambiar un atributo de la capa. Medimos asignaciones descartadas sobre cuatro mil noventa y seis, y tokens que se quedaron sin ningún experto. Sin límite no se descarta nada, y los errores son los de siempre: cero coma cero uno seis y cero coma cero uno ocho siete. El colapsado, con cupo, es un desastre: con factor uno pierde el cincuenta y cinco por ciento de las asignaciones, y el cuarenta y uno por ciento de los tokens no ve ningún experto. Su error salta a cero coma cincuenta y ocho. El balanceado, con factor uno, descarta el dos coma cuatro por ciento; con uno coma veinticinco, nada.

pct = lambda v: " ".join(f"{100 * x:4.1f}" for x in v)
for alfa in ("0.0", "0.1"):                  # sin y con balanceo (ep. 4)
    red, desc, sin, mse = cargar(alfa), [], [], []
    for factor in (None, 1.0, 1.25, 2.0):         # None: sin límite
        red.capa.factor_capacidad = factor
        with torch.no_grad():
            mse.append(F.mse_loss(red(Xt), Yt).item())
        desc.append(red.capa.ultimo["descartadas"] / 4096)
        sin.append(red.capa.ultimo["tokens_sin_experto"] / 2048)
    print(f"α={alfa} % descart. {pct(desc)} | % sin experto {pct(sin)}")
    print(f"      MSE {' '.join(f'{m:.4f}' for m in mse)}")

α=0.0 % descart.  0.0 55.2 47.5 28.8 | % sin experto  0.0 41.3 33.7 18.5
      MSE 0.0162 0.5803 0.4985 0.3024
α=0.1 % descart.  0.0  2.4  0.0  0.0 | % sin experto  0.0  1.2  0.0  0.0
      MSE 0.0187 0.0338 0.0187 0.0187

Tokens sin ningún experto (salida real)

Las mismas cifras en barras. A la izquierda, el modelo que colapsó: aun duplicando la capacidad, casi uno de cada cinco tokens pasa de largo, porque sus favoritos reciben más del doble de lo que entra. A la derecha, el balanceado: con un margen de veinticinco por ciento ya no pierde nada. El capacity factor no arregla un router desbalanceado; solo funciona encima de uno que reparte bien.

• sin balanceo · 1,0: 41,3 %
• sin balanceo · 1,25: 33,7 %
• sin balanceo · 2,0: 18,5 %
• α=0,1 · 1,0: 1,2 %
• α=0,1 · 1,25: 0 %
• α=0,1 · 2,0: 0 %

Error de test según el factor

El efecto en el error de test, medido como MSE, mean squared error: el promedio de los errores al cuadrado. Sin balanceo, el error baja a medida que agrandamos el cupo, pero ni con factor dos se acerca a lo que daba sin límite. El balanceado casi no se entera: el error con factor uno coma veinticinco ya es igual al de sin límite. Como referencia, sin cupo daban cero coma cero uno seis y cero coma cero uno ocho siete: la curva verde, con factor uno coma veinticinco, ya cae exactamente ahí.

El precio del margen: relleno

Y el otro lado de la balanza: el relleno. La carga de cada experto en el test. Capacidad con factor dos: mil veinticuatro lugares cada uno. Lugares ocupados: la carga de cada experto, recortada en C, sumada. El balanceado usa exactamente la mitad de los lugares: con factor dos, la mitad de las cuentas de los expertos se hacen sobre ceros. El colapsado usa todavía menos, el treinta y seis por ciento, y además pierde tokens, porque su experto favorito recibe mil setecientas cuarenta y cuatro asignaciones.

for alfa in ("0.0", "0.1"):
    red = Red(16, 32, 8, MoE(32, 32, E=8, k=2))
    red.load_state_dict(torch.load(f"{AQUI}/../datos/ep04_alfa{alfa}.pt"))
    with torch.no_grad():
        red(Xt)
    carga = red.capa.ultimo["carga"]               # asignaciones por experto
    C = math.ceil(2.0 * 2048 * 2 / 8)              # factor 2,0 → 1024
    usados = carga.clamp(max=C).sum().item()       # lugares con un token
    print(f"α={alfa}: carga máx {carga.max().item():>4} | lugares usados "
          f"{usados}/{8 * C} = {usados / (8 * C):.0%}")

α=0.0: carga máx 1744 | lugares usados 2917/8192 = 36%
α=0.1: carga máx  545 | lugares usados 4096/8192 = 50%

El compromiso

• Factor bajo: se descartan tokens
• Factor alto: cuentas y memoria sobre relleno
• El balanceo es lo que permite un factor bajo
• Ojo: acá el cupo se puso solo al evaluar

Las dos puntas. Con un factor bajo se aprovecha cada lugar, pero se descartan tokens, y un token descartado se saltea toda la capa. Con un factor alto no se pierde nada, pero se calcula y se guarda relleno. Switch reporta que su modelo funcionaba mejor con factores bajos, como uno o uno coma veinticinco. Y lo que hace posible un factor bajo es un router bien balanceado: por eso estos dos episodios van juntos. Una salvedad del experimento: nuestros modelos se entrenaron sin cupo y se lo pusimos recién al evaluar. Entrenando con cupo, el modelo aprende sabiendo que algunos tokens se descartan.

Parte 4 · Sin descartar nada

Hay otra salida: no fijar la forma, y escribir kernels que sepan trabajar con expertos de cargas distintas.

Gale et al., 2022 — MegaBlocks: Efficient Sparse Training with Mixture-of-Experts · arXiv:2211.15841

hasta 40 % · «Our approach never drops tokens and maps efficiently to modern hardware, enabling end-to-end training speedups of up to 40% over MoEs trained with the state-of-the-art Tutel library».

MegaBlocks, de Stanford y otros, es la referencia del enfoque dropless, sin descartes. Reescriben la capa MoE como operaciones con matrices ralas por bloques, block-sparse en inglés, y escriben kernels de GPU, graphics processing unit, para eso. Su sistema nunca descarta tokens, y reportan hasta un cuarenta por ciento de aceleración de entrenamiento de punta a punta contra MoE entrenados con Tutel, otra librería del momento.

Dos formas de resolver la carga variable

Las dos familias, lado a lado. Con capacidad: un buffer fijo, los tokens que no entran siguen por la residual, y el margen se paga con relleno. A cambio, todo tiene forma conocida y es fácil de compilar. Dropless: cada experto procesa exactamente los tokens que le tocaron, sin perder ni rellenar nada, pero hace falta código de bajo nivel que sepa multiplicar bloques de tamaño variable. Nuestra capa de juguete, sin factor, es dropless: un loop por experto con su cantidad exacta. En el procesador de una notebook eso es fácil; lo difícil es hacerlo rápido en una placa de video.

Resumiendo

• Hardware y compiladores piden formas fijas
• $C = \lceil \text{factor} \cdot Tk/E \rceil$ lugares por experto
• Overflow: el token sigue solo por la residual
• Colapsado + factor 1,0: 41 % de tokens sin experto
• Balanceado + factor 1,25: cero descartes
• Próximo: top-1, Switch y la estabilidad del router

Papers del episodio

• Fedus, Zoph & Shazeer, 2021 — Switch Transformers
• Lepikhin et al., 2020 — GShard
• Gale et al., 2022 — MegaBlocks: Efficient Sparse Training with Mixture-of-Experts