Switch y la estabilidad

Mixture of Experts, episodio seis: Switch y la estabilidad. Ya tenemos un MoE, mixture of experts, con router, balanceo y capacidad. Hoy, dos ideas que lo hicieron práctico a gran escala: usar un solo experto por token, y evitar que los números del router se rompan cuando entrenás con poca precisión.

El plan

• Switch Transformer: top-1 en lugar de top-2
• El truco para que el router siga aprendiendo con k = 1
• bfloat16: pocos bits, redondeos grandes
• Router en float32 y router z-loss

Cuatro piezas. Switch Transformer, de Google, mostró que alcanza con mandar cada token a un solo experto. Pero con k igual a uno hay una trampa en el gradiente, y la vamos a ver en código. Después, la precisión: entrenar en bfloat16, un formato de dieciséis bits, es rápido, pero redondea fuerte. Y las dos curas: calcular el router en float32, y una pérdida extra que mantiene chicos los logits, la router z-loss.

Parte 1 · Top-1 alcanza

Hasta ahora cada token iba a dos expertos. ¿Qué pasa si va a uno solo?

La capa Switch

el router de siempre: E probabilidades

un solo experto por token

se multiplica por p, sin renormalizar

La capa Switch, paso a paso. El router es el de siempre: W sub r, una matriz de d por E, multiplica al vector x del token, y un softmax da E probabilidades que suman uno. En el dibujo, cuatro expertos de ejemplo. i asterisco es el índice de la probabilidad más alta: el único experto que se va a calcular. Y la salida es la salida de ese experto multiplicada por su probabilidad, p sub i asterisco. En el dibujo, solo la columna del experto cero entra en la suma, pesada por cero coma cincuenta y dos; las otras tres ni se calculan. Ese factor p parece un detalle, pero es lo único que conecta al router con la pérdida.

Si renormalizás, el router no aprende

Probémoslo con el MoE de juguete de la serie. Una capa con ocho expertos y k igual a uno, una vez renormalizando los pesos del top-k para que sumen uno, como hacíamos con top-2, y otra vez sin renormalizar. Hacemos backward de una pérdida cualquiera. Miramos los pesos de tres tokens y la norma del gradiente que llega a W sub r. Renormalizando, un solo peso dividido por sí mismo da uno: los pesos son uno coma cero, una constante, y el gradiente del router es exactamente cero. El router nunca aprendería. Sin renormalizar, los pesos son las probabilidades, cero coma dos tres, cero coma tres uno, y el gradiente llega.

for renormalizar in [True, False]:
    torch.manual_seed(1)
    moe = MoE(d=16, h=32, E=8, k=1, renormalizar=renormalizar)   # top-1
    moe(x).pow(2).mean().backward()            # una pérdida cualquiera
    g = moe.ultimo["pesos"][:3, 0].tolist()
    print(f"renormalizar={renormalizar!s:5}  pesos {[round(v, 3) for v in g]}"
          f"  |∇W_r| = {moe.router.weight.grad.norm():.4f}")

renormalizar=True   pesos [1.0, 1.0, 1.0]  |∇W_r| = 0.0000
renormalizar=False  pesos [0.234, 0.307, 0.232]  |∇W_r| = 0.0051

Lo que se gana con k = 1, según el paper

• Menos cómputo en el router: un experto por token
• Capacidad por experto al menos a la mitad
• Implementación más simple y menos comunicación

El paper de Switch lista tres ganancias. El router hace menos trabajo, porque cada token va a un solo experto. La capacidad de cada experto, cuántos tokens procesa por batch, que vimos en el episodio cinco, se puede reducir al menos a la mitad. Y la implementación es más simple, con menos comunicación: cada token viaja a un solo lugar. Pierde la mezcla de dos opiniones por token. Midámoslo.

Top-2 contra top-1 en la tarea sintética

Dos configuraciones: top-2 renormalizado, y top-1 al estilo Switch, sin renormalizar. Mismos ocho expertos, misma semilla. Las dos entrenan ochocientos pasos con la pérdida de balanceo de Switch, alfa cero coma cero uno. Y contamos las asignaciones en un batch de doscientos cincuenta y seis tokens. Top-2 llega a cero coma doscientos ocho de error cuadrático medio en test, el MSE, mean squared error. Top-1, a cero coma doscientos diecinueve: un poco peor, con la mitad de asignaciones, doscientas cincuenta y seis contra quinientas doce. Lo repetimos aparte con dos semillas más, y top-1 siempre quedó algo peor. Cada asignación es un experto que corre y un vector que viaja.

for k, renorm in [(2, True), (1, False)]:      # top-2 clásico y Switch
    torch.manual_seed(0)
    moe = MoE(d=32, h=64, E=8, k=k, renormalizar=renorm)
    red = Red(16, 32, 8, moe)
    entrenar(red, X, Y, pasos=800, alfa=0.01)  # con L_aux de Switch
    with torch.no_grad():
        mse = F.mse_loss(red(Xt), Yt).item()
        red(Xt[:256])                           # un batch de 256 tokens
    n = moe.ultimo['carga'].sum().item()
    print(f"top-{k}: MSE {mse:.3f}, {n} asignaciones")

top-2: MSE 0.208, 512 asignaciones
top-1: MSE 0.219, 256 asignaciones

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

7× más rápido · Su Switch-Base de 64 expertos «achieves the same quality in one-seventh the time of the T5-Base», con los mismos FLOPs por ejemplo.

La cifra que hizo famoso al paper. Con el mismo cómputo por token, medido en FLOPs, floating point operations, las cuentas que hace la red, su Switch-Base de sesenta y cuatro expertos llegó a la misma calidad que T5-Base, un transformer denso de Google, en un séptimo del tiempo. Más parámetros, sin más cuentas por token.

Parte 2 · Pocos bits, redondeos grandes

Para ir rápido, los modelos grandes se entrenan en bfloat16. Y el router es justo la parte que peor lo tolera.

float32 contra bfloat16

Dos formatos de punto flotante. float32 usa ocho bits para el exponente, que fija el rango, y veintitrés para la mantisa, que fija la precisión. Entre treinta y dos y sesenta y cuatro, los números representables están separados por dos a la menos dieciocho. bfloat16 conserva los ocho bits de exponente, así que llega igual de lejos, pero solo tiene siete de mantisa. Entre treinta y dos y sesenta y cuatro, los números van de a cero coma veinticinco. Son dieciséis bits menos de mantisa: los pasos son dos a la dieciséis veces más grandes.

Por qué el router sufre

q: el exponente; el paso crece con el número

redondear = sumar un error δ

la exponencial multiplica el error por e^ℓ

La cuenta de por qué duele. En bfloat16, si x está entre dos a la q y dos a la q más uno, con q un entero, el exponente, el paso entre números vecinos es dos a la q menos siete. Cuanto más grande el número, más grueso el paso. Redondear un logit ele es sumarle un error delta, de hasta medio paso. Y el softmax eleva e, el número de Euler, a la ele: el error queda multiplicado. e a la ele más delta es e a la ele por e a la delta, que con delta chico es e a la ele por uno más delta. En el dibujo, la pendiente de la exponencial: en ele igual a dos es siete coma cuatro, siete veces más empinada que en cero. Logits grandes, probabilidades torcidas.

Redondear logits en bfloat16

Primero, tres números redondeados a bfloat16. Después, una función que calcula el peso del primer experto entre dos, redondeando los logits al formato que le digamos antes del softmax. Probamos la misma diferencia de cero coma dos entre logits, a tres escalas. Pi queda en tres coma uno cuatro; cuarenta coma tres, en cuarenta coma veinticinco; mil quinientos coma siete, en mil quinientos cuatro. Y el peso: con logits chicos, cero coma cincuenta y cinco en los dos formatos. Con logits cerca de cuarenta, bfloat16 da cero coma quinientos sesenta y dos. Cerca de trescientos, los dos logits caen en el mismo número y el peso queda en cero coma cinco: un empate que no existía. Ese ejemplo es nuestro; el paper de ST-MoE dice que el redondeo puede cambiar mucho la probabilidad que escala la salida del experto.

for x in [3.14159, 40.3, 1500.7]:
    print(f"{x:>8} → bfloat16 {torch.tensor(x).bfloat16().item()}")

def gate(logits, dtype):                  # peso del primer experto de dos
    l = torch.tensor(logits).to(dtype).float()   # redondear y volver
    return l.softmax(-1)[0].item()

for l in ([0.3, 0.1], [40.3, 40.1], [300.3, 300.1]):
    print(f"{l[0]:>5}/{l[1]}: fp32 {gate(l, torch.float32):.3f}"
          f" · bf16 {gate(l, torch.bfloat16):.3f}")

 3.14159 → bfloat16 3.140625
    40.3 → bfloat16 40.25
  1500.7 → bfloat16 1504.0
  0.3/0.1: fp32 0.550 · bf16 0.550
 40.3/40.1: fp32 0.550 · bf16 0.562
300.3/300.1: fp32 0.550 · bf16 0.500

Parte 3 · Las dos curas

Calcular el router con más bits, y pedirle al router que mantenga chicos sus logits.

Selective precision, en Granite

La primera cura se llama selective precision: todo el modelo en bfloat16, salvo el router, que trabaja en float32. Leemos el código fuente del router de Granite, el MoE abierto de IBM que usamos en esta serie, tal como viene en transformers. Filtramos las dos líneas que importan. Sin comentarios. Los logits salen de la capa lineal en el tipo del modelo, y recién ahí se pasan a float32 con punto float. El top-k y el softmax corren en float32, y los pesos vuelven al tipo del modelo con type as. Ojo: el redondeo de los logits ya ocurrió; lo que se protege es la exponencial. Switch hacía algo más fuerte: pasaba a float32 la entrada del router.

fuente = inspect.getsource(gm.GraniteMoeTopKRouter.forward)
for linea in fuente.splitlines():
    if "F.linear" in linea or "softmax" in linea:
        print(linea.strip().split("  #")[0])

router_logits = F.linear(hidden_states, self.weight).float()
top_k_weights = torch.softmax(top_k_logits, dim=-1).type_as(hidden_states)

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

bf16: divergió · Switch-Base en bfloat16 «diverged»; con selective precision, calidad de float32 (−1,716 contra −1,718) y la velocidad de bfloat16.

El experimento del paper. Entrenaron el mismo modelo de tres maneras. Todo en float32: estable, pero más lento. Todo en bfloat16: divergió, la pérdida se fue a cualquier lado. Y con el router en float32 y el resto en bfloat16: la calidad de float32, con casi la velocidad de bfloat16.

La router z-loss

log-sum-exp: un máximo suave

B tokens, N expertos

empuja hacia abajo, más fuerte al más probable

tarea + balanceo + z-loss

La segunda cura. LSE, log-sum-exp, es el logaritmo de la suma de las exponenciales de los logits: un máximo suave. En el dibujo, los logits de un ejemplo: cuatro, uno, cero coma cinco, menos uno. LSE da cuatro coma cero ocho, apenas más que el máximo. La z-loss es el promedio, sobre los B tokens del batch, de LSE al cuadrado. Penaliza logits grandes. Derivando con la regla de la cadena: dos sobre B, por LSE, por la probabilidad de cada experto. En el dibujo, ese softmax: el gradiente tira hacia abajo de todos los logits, más fuerte del más probable. Se suma a la pérdida de la tarea y a la de balanceo, con un coeficiente c sub z de cero coma cero cero uno.

La z-loss en tres líneas

En código, con dos tokens y cuatro expertos. torch punto logsumexp, por fila. Al cuadrado y promediado: la z-loss. Comparamos el gradiente de autograd con la fórmula. Y al final, restamos tres a todos los logits. LSE da cuatro coma cero ocho y dos coma ocho dos, z da doce coma treinta y uno, y autograd coincide con la fórmula hasta dos por diez a la menos siete. El gradiente más grande, tres coma siete seis, es el del logit más alto. Y el detalle clave: después de restar tres, el softmax no cambia, pero z baja a cero coma seis. La z-loss puede achicar los logits sin tocar ninguna decisión del router.

Cada línea, su término:

• lse = torch.logsumexp(logits, dim=-1) → $\mathrm{LSE}(x^{(i)})$ — una por token
• z = (lse ** 2).mean() → $L_z = \tfrac{1}{B}\textstyle\sum_i \mathrm{LSE}^2$ — promedio de cuadrados
• analitico = 2 / B → $\tfrac{2}{B}\,\mathrm{LSE}\; p_j$ — el gradiente de la fórmula
• corrido = logits.detach() - 3 → $\mathrm{softmax}(x - c) = \mathrm{softmax}(x)$ — restar no cambia el ruteo

logits = torch.tensor([[4.0, 1.0, 0.5, -1.0],    # 2 tokens, 4 expertos
                       [2.0, 2.0, 0.0, 0.0]], requires_grad=True)
lse = torch.logsumexp(logits, dim=-1)            # log Σ_j e^{x_j}
z = (lse ** 2).mean()                            # router z-loss
z.backward()
print("LSE:", [round(v, 3) for v in lse.tolist()], "· z =", f"{z:.3f}")

B = logits.shape[0]
analitico = 2 / B * lse[:, None] * logits.softmax(-1)   # (2/B)·LSE·p
print("máx |grad − fórmula| =", f"{(logits.grad - analitico).abs().max():.1e}")
print("grad fila 0:", [round(v, 3) for v in logits.grad[0].tolist()])

corrido = logits.detach() - 3                     # restar 3 a todos
print("misma softmax:", torch.allclose(corrido.softmax(-1), logits.softmax(-1)))
print("z corrido =", f"{(torch.logsumexp(corrido, -1) ** 2).mean():.3f}")

LSE: [4.083, 2.82] · z = 12.313
máx |grad − fórmula| = 2.4e-07
grad fila 0: [3.757, 0.187, 0.113, 0.025]
misma softmax: True
z corrido = 0.603

Entrenar con y sin z-loss

Ahora en entrenamiento. El mismo MoE de top-2, ochocientos pasos. En cada paso calculamos z sobre los logits del router. La sumamos a la pérdida con su coeficiente. Sin z-loss y con el valor del paper. Sin z-loss, z llega a veintiocho coma seis y el logit ganador promedio es cinco coma uno. Con c sub z de cero coma cero cero uno, z queda en cuatro coma cinco y el logit ganador en uno coma siete. El MSE pasa de cero coma doscientos ocho a cero coma doscientos veintiocho, pero con otras semillas esa diferencia va y viene: está dentro del ruido. Ojo: con ocho expertos y logits de cinco, este modelo nunca corrió peligro. Lo que se ve es el mecanismo: los logits bajan.

def entrenar_z(c_z, pasos=800):
    torch.manual_seed(0)
    moe = MoE(d=32, h=64, E=8, k=2); red = Red(16, 32, 8, moe)
    opt = torch.optim.Adam(red.parameters(), lr=3e-3)
    g = torch.Generator().manual_seed(0)
    for _ in range(pasos):
        i = torch.randint(0, len(X), (256,), generator=g)
        mse = F.mse_loss(red(X[i]), Y[i])
        z = (torch.logsumexp(moe.ultimo["logits"], -1) ** 2).mean()
        loss = mse + 0.01 * moe.ultimo["l_aux"] + c_z * z
        opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad():
        mse = F.mse_loss(red(Xt), Yt).item(); lg = moe.ultimo["logits"]
    print(f"c_z={c_z}: MSE {mse:.3f} z {(lg.logsumexp(-1) ** 2).mean():.2f}"
          f" máx {lg.max(-1).values.mean():.2f}")
for c_z in [0.0, 0.001]:
    entrenar_z(c_z)

c_z=0.0: MSE 0.208 z 28.65 máx 5.10
c_z=0.001: MSE 0.228 z 4.53 máx 1.71

Lo que cambia la z-loss (tarea sintética)

En barras. La z sin penalizar, y con penalización: seis veces menos. El logit ganador sin z-loss, y con z-loss, cerca de uno. Chico es preciso: cerca de uno, bfloat16 da pasos de menos de una centésima; cerca de trescientos, de dos unidades.

• z, sin z-loss: 28,65
• z, c_z = 0,001: 4,53
• logit máx., sin: 5,10
• logit máx., con: 1,71

Zoph et al., 2022 — ST-MoE: Designing Stable and Transferable Sparse Expert Models · arXiv:2202.08906

3/3 estables · Sin la router z-loss, 4 de 6 corridas estables; con ella, 3 de 3, y «without quality degradation». Eligen c_z = 0,001.

El paper que la propuso. En ST-MoE, de Google, el modelo base fue estable en cuatro de seis corridas; con la router z-loss, en tres de tres, sin perder calidad, incluso con una leve mejora. Y con la z-loss escalaron a un MoE de doscientos sesenta y nueve mil millones de parámetros.

La receta que quedó

• k = 1 o 2: top-1 sin renormalizar, o top-2 renormalizado
• Softmax del router en float32, el resto en bfloat16
• Router z-loss con c_z ≈ 0,001
• Medir siempre: z, logit máximo, carga por experto

Lo que quedó como práctica. Uno o dos expertos por token; con uno, no renormalices, o el router se queda sin gradiente. Todo lo que pasa por una exponencial, en float32. Una z-loss chica, del orden de cero coma cero cero uno. Y monitorear la z, el logit máximo y la carga por experto: avisan antes de que algo se rompa.

Resumiendo

• Switch: un experto por token, y = p·E(x)
• Con k = 1, renormalizar mata el gradiente del router
• bfloat16: 7 bits de mantisa, pasos que crecen con el número
• Selective precision: router en float32
• z-loss = media de LSE²: logits chicos, mismo ruteo
• Próximo: expertos finos y compartidos (DeepSeekMoE)

Papers del episodio

• Fedus, Zoph y Shazeer, 2021 — Switch Transformers
• Zoph et al., 2022 — ST-MoE: Designing Stable and Transferable Sparse Expert Models