Balancear sin pérdida auxiliar

Mixture of Experts, episodio ocho: balancear sin pérdida auxiliar. En el episodio cuatro vimos que el router de un MoE, mixture of experts, colapsa si nadie lo frena, y lo frenamos con una pérdida de balanceo. Hoy, otra forma: un número por experto, que se ajusta solo y no toca el gradiente. Es lo que usa DeepSeek-V3.

El plan

• El problema: la pérdida auxiliar compite con la tarea
• Un sesgo b_i que solo cambia a quién se elige
• La regla: bajarlo si sobra carga, subirlo si falta
• Medir con MaxVio contra L_aux en el juguete

Cuatro pasos. Primero, por qué la pérdida auxiliar es un compromiso incómodo. Después, la idea: sumarle a cada experto un sesgo b sub i, pero solo para decidir quién entra en el top-k, no para pesar su salida. La regla para ajustarlo es casi un termostato. Y lo comparamos en nuestro MoE de juguete con una métrica nueva, MaxVio, la violación máxima de carga.

Parte 1 · Dos objetivos, un gradiente

La pérdida de balanceo no es gratis: su gradiente se mezcla con el de la tarea.

La pérdida auxiliar, de nuevo

f: fracción de asignaciones; P: probabilidad media

α elige cuánto pesa el balanceo

los dos gradientes se suman en los mismos pesos θ

Recordemos la pérdida del episodio cuatro. E es la cantidad de expertos, f sub i la fracción de asignaciones que recibió el experto i, y P sub i la probabilidad media que el router le dio. Vale uno cuando todo es parejo. En el dibujo, una carga despareja de ejemplo. Se suma a la pérdida de la tarea multiplicada por alfa. Y acá está el problema: el gradiente total es la suma de los dos gradientes, y los dos empujan los mismos pesos theta. En el dibujo, el gradiente de la tarea más alfa por el del balanceo da una dirección que ya no es la que la tarea pedía. Con alfa chico, la carga queda despareja; con alfa grande, la tarea empeora.

El compromiso de α

Un esquema en dos dimensiones, con vectores inventados. La flecha celeste es hacia dónde querría moverse la red para resolver la tarea. La ámbar es lo que pide el balanceo, multiplicado por alfa. Y la violeta, la suma de las dos, es lo que de verdad pasa. Cada paso de entrenamiento paga un poco de tarea para comprar balance. La pregunta del paper que vemos hoy es si se puede comprar balance sin pagar en el gradiente.

Wang et al., 2024 — Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts · arXiv:2408.15664

MaxVio 0,72 → 0,04 · En el modelo de 1B: con pérdida auxiliar, perplejidad 9,56 y MaxVio global 0,72; Loss-Free, 9,50 y 0,04. Probado en modelos de hasta 3B y 200B tokens.

El paper. Wang y coautores, de DeepSeek, comparan la pérdida auxiliar contra su método sin pérdida en modelos de hasta tres mil millones de parámetros. En el de mil millones, la pérdida auxiliar deja una violación máxima de carga de cero coma setenta y dos, y su método, de cero coma cero cuatro. Y la perplejidad, que mide qué tan bien predice el texto, también mejora un poco: de nueve coma cincuenta y seis a nueve coma cincuenta.

Parte 2 · Un sesgo para elegir

La idea entera cabe en una línea: el sesgo cambia quién entra en el top-k, pero no cuánto pesa.

Elegir con sesgo, pesar sin sesgo

s: puntaje del router; b: un número por experto

el peso sale del puntaje original

se ajusta con una regla, no con backprop

Con números. Un token, cuatro expertos, logits del router dos, uno coma cinco, uno coma cuatro y cero coma dos. Sin sesgo, ganan el cero y el uno. Ahora sumamos un sesgo por experto: menos cero coma seis al cero, que está sobrecargado, y más cero coma tres al dos, que está ocioso. S es el conjunto que gana con s más b. En el dibujo, las sumas: entran el dos y el uno. Pero el peso de cada elegido sale de las probabilidades originales, sin el sesgo. En el dibujo, el softmax de los logits sin tocar: el experto dos entra con cero coma doscientos treinta y seis. Y b no participa de ninguna cuenta que tenga gradiente: no se entrena con backprop, se ajusta con una regla aparte.

Elegir con s + b, pesar con s

A mano, en dos líneas; el router de juguete de la serie hace lo mismo por dentro. El sesgo del ejemplo. Para elegir, top-k sobre los logits más el sesgo. Para pesar, las probabilidades originales de los elegidos, sin renormalizar, así se ve de dónde sale cada peso. E imprimimos esas probabilidades. Sin sesgo, elige cero y uno, con cero coma cuatro tres y cero coma dos seis. Con sesgo, elige dos y uno, y los pesos son cero coma dos tres seis y cero coma dos seis uno: exactamente las probabilidades originales de esos expertos. El sesgo movió la decisión y no tocó ningún peso.

logits = torch.tensor([2.0, 1.5, 1.4, 0.2])      # 1 token, 4 expertos
sesgo = torch.tensor([-0.6, 0.0, 0.3, 0.0])       # e0 sobrecargado, e2 ocioso
probs = logits.softmax(-1)                        # puntaje original

for b in [torch.zeros(4), sesgo]:
    idx = (logits + b).topk(2).indices            # ELEGIR con s + b
    pesos = probs[idx]                            # PESAR con s, sin b
    print("con b:" if b.any() else "sin b:", idx.tolist(), "→",
          [round(p, 3) for p in pesos.tolist()])
print("probs originales:")
print([round(p, 3) for p in probs.tolist()])

sin b: [0, 1] → [0.431, 0.261]
con b: [2, 1] → [0.236, 0.261]
probs originales:
[0.431, 0.261, 0.236, 0.071]

La regla y la métrica

c_i: asignaciones del experto i en el batch

sobrecargado: baja · ocioso: sube

0 = parejo; 1 = alguien con el doble de lo justo

La regla. c sub i es cuántas asignaciones recibió el experto i en el batch, y c barra, el promedio: T tokens por k elegidos, dividido E. Después de cada paso, el sesgo de cada experto se mueve una cantidad fija u: baja si el experto tuvo más que el promedio, sube si tuvo menos. sign devuelve más uno o menos uno, según el signo. Es un termostato. En el dibujo, la carga menos el promedio, del batch real del próximo snippet. Para medir el balance, el paper define MaxVio: cuánto se pasa el experto más cargado respecto del promedio, en proporción. Cero es perfecto; uno es que alguien recibe el doble de lo justo.

Medir MaxVio

La métrica en código. Máximo menos media, dividido la media. Tres cargas de ejemplo, todas con quinientas doce asignaciones: doscientos cincuenta y seis tokens por dos. La pareja da cero. La despareja, con un experto en noventa contra un promedio de sesenta y cuatro, da cero coma cuarenta y uno. Y el colapso, con un experto que se lleva cuatrocientas, da cinco coma veinticinco.

def maxvio(carga):                      # (máx_i c_i − c̄) / c̄
    c = carga.float()
    return ((c.max() - c.mean()) / c.mean()).item()

ejemplos = {"pareja":    [64, 64, 64, 64, 64, 64, 64, 64],
            "despareja": [90, 40, 70, 64, 30, 80, 60, 78],
            "colapso":   [400, 40, 30, 20, 10, 6, 4, 2]}
for nombre, c in ejemplos.items():
    print(f"{nombre:10} suma {sum(c)} · MaxVio {maxvio(torch.tensor(c)):.2f}")

pareja     suma 512 · MaxVio 0.00
despareja  suma 512 · MaxVio 0.41
colapso    suma 512 · MaxVio 5.25

Un paso de la regla

Un paso, con un batch de verdad por el MoE sin entrenar. La velocidad u, cero coma cero uno. La carga de cada experto en este batch. El error: promedio menos carga. Y la actualización, una línea. Por último, preguntamos si el sesgo tiene gradiente. El experto cinco recibió ciento treinta asignaciones contra un promedio de sesenta y cuatro: su signo es menos uno y su sesgo baja a menos cero coma cero uno. El seis recibió veintinueve: sube. Y requires grad es falso: el sesgo es un buffer, vive fuera del gradiente.

u = 0.01                                     # velocidad de ajuste del sesgo
with torch.no_grad():
    red(X)                                   # un batch: 256 tokens, top-2
c = moe.ultimo["carga"].float()              # asignaciones por experto
error = c.mean() - c                         # > 0: le faltan tokens
moe.sesgo += u * torch.sign(error)           # b_i ← b_i + u·sign(c̄ − c_i)
print("carga  :", " ".join(str(x) for x in c.int().tolist()))
print("media  :", c.mean().item())
print("signo  :", " ".join("+" if e > 0 else "−" for e in error))
print("sesgo/u:", " ".join(f"{b / u:+.0f}" for b in moe.sesgo.tolist()))
print("¿gradiente?", moe.sesgo.requires_grad)

carga  : 55 83 56 94 30 130 29 35
media  : 64.0
signo  : + − + − + − + +
sesgo/u: +1 -1 +1 -1 +1 -1 +1 +1
¿gradiente? False

El termostato

El ciclo completo. El router elige con el puntaje más el sesgo. Se cuenta la carga de cada experto en el batch. Se compara con el promedio. Y se corrige el sesgo un paso fijo. Vuelta a empezar en el próximo batch. Mientras tanto, el gradiente de la tarea entrena todo lo demás sin ninguna interferencia.

Parte 3 · Comparar en el juguete

Tres maneras de entrenar el mismo MoE: sin balanceo, con pérdida auxiliar, y con el sesgo.

Nada, L_aux o sesgo

Una sola función de entrenamiento, ochocientos pasos, ocho expertos, top-2. La pérdida auxiliar entra solo si alfa no es cero. El sesgo se ajusta después de cada paso, solo si u no es cero. Tres corridas. Sin balanceo: error cuadrático medio, el MSE, mean squared error, de cero coma doscientos veinte, y MaxVio en test de cero coma cincuenta y tres. Con la pérdida auxiliar: cero coma doscientos ocho y cero coma treinta y siete. Con el sesgo: cero coma doscientos quince, y MaxVio de cero coma cero nueve. El sesgo balancea cuatro veces mejor, pero el menor error es el de la pérdida auxiliar. Con dos semillas más, probadas aparte, se repite el orden. En este juguete la pérdida auxiliar no estorbó: el argumento de la parte uno viene del paper, a escala, y del episodio cuatro.

def entrenar(alfa=0.0, u=0.0, pasos=800):
    torch.manual_seed(0)
    moe = MoE(32, 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); mv = []
    for _ in range(pasos):
        i = torch.randint(0, len(X), (256,), generator=g)
        loss = F.mse_loss(red(X[i]), Y[i]) + alfa * moe.ultimo["l_aux"]
        opt.zero_grad(); loss.backward(); opt.step()
        c = moe.ultimo["carga"].float(); mv.append(maxvio(c))
        moe.sesgo += u * torch.sign(c.mean() - c)    # 0 si u = 0
    with torch.no_grad():
        mse = F.mse_loss(red(Xt), Yt).item()
    return mse, maxvio(moe.ultimo["carga"]), mv

for nombre, kw in {"nada": {}, "L_aux": dict(alfa=0.01),
                   "sesgo": dict(u=0.01)}.items():
    mse, mv_test, mv = entrenar(**kw); historia[nombre] = mv
    print(f"{nombre:6} MSE {mse:.3f} · MaxVio test {mv_test:.2f}")

nada   MSE 0.220 · MaxVio test 0.53
L_aux  MSE 0.208 · MaxVio test 0.37
sesgo  MSE 0.215 · MaxVio test 0.09

MaxVio durante el entrenamiento

Y durante el entrenamiento, promediando de a cincuenta pasos. Sin balanceo, la violación arranca cerca de uno y se queda arriba de cero coma cinco. La pérdida auxiliar la baja, pero oscila entre cero coma treinta y cinco y cero coma cincuenta y cinco. El sesgo la lleva a cerca de cero coma dos en los primeros cien pasos, y ahí se queda. Ojo: en el batch, la violación nunca llega a cero, porque cada batch trae tokens distintos y la regla corrige con un paso de retraso.

Parte 4 · La velocidad del termostato

La regla tiene un solo número para elegir: u, cuánto se mueve el sesgo en cada paso.

Tres velocidades

Probamos tres valores de u con la misma función de antes. Cero coma cero cero uno, el del paper; cero coma cero uno, el que usamos; y cero coma uno. Para ver la velocidad, promediamos MaxVio en los primeros cien pasos y en los últimos cien. Con cero coma cero cero uno, la violación baja de cero coma noventa y uno a cero coma cuarenta y seis: el termostato es demasiado lento para ochocientos pasos. Con cero coma cero uno, de cero coma cuarenta a cero coma dieciocho. Con cero coma uno, arranca más pareja, en cero coma veintiocho, y termina en cero coma diecinueve, con el menor error de los tres. Con dos semillas más se repite; lo que queda es que u depende de cuántos pasos entrenás. Según el paper, un u muy alto empeora el balance más adelante.

for u in [0.001, 0.01, 0.1]:              # velocidad del termostato
    mse, mv_test, mv = entrenar(u=u)
    primeros = sum(mv[:100]) / 100            # MaxVio medio, pasos 1–100
    ultimos = sum(mv[-100:]) / 100            # MaxVio medio, pasos 701–800
    print(f"u={u:<5} MSE {mse:.3f} · MaxVio {primeros:.2f} → {ultimos:.2f}")

u=0.001 MSE 0.209 · MaxVio 0.91 → 0.46
u=0.01  MSE 0.215 · MaxVio 0.40 → 0.18
u=0.1   MSE 0.199 · MaxVio 0.28 → 0.19

Letra chica

• u depende del largo del entrenamiento: 0,001 en el paper
• Balance no es calidad: acá el error quedó entre los otros dos
• Se balancea por batch: con batch chico, más ruido
• Nuestro sesgo suma a los logits; DeepSeek-V3 a puntajes sigmoide

Detalles que importan. La velocidad u depende de cuánto vas a entrenar. En el paper, cero coma cero cero uno funcionó mejor, pero con cien mil millones de tokens; en nuestros ochocientos pasos, se queda corto. Balancear no garantiza mejor calidad: en este juguete, el error quedó entre la pérdida auxiliar y nada. La carga se mide por batch, así que con batches chicos la regla persigue ruido. Y nuestro router suma el sesgo a los logits del softmax; DeepSeek-V3 calcula los puntajes con una sigmoide, una función que lleva cada número a algo entre cero y uno, y suma el sesgo ahí.

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

γ = 0,001 · «The bias term is only used for routing»; γ = 0,001 durante los primeros 14,3T tokens y 0,0 en los últimos 500B. Igual suman una pérdida por secuencia con α = 0,0001.

Así se usa a escala. DeepSeek-V3 lo dice textual: el sesgo solo se usa para rutear, y el peso que multiplica la salida del experto sale del puntaje original. Lo bajan o suben una cantidad gama, cero coma cero cero uno, durante los primeros catorce coma tres billones de tokens, y lo congelan en los últimos quinientos mil millones. Y un detalle honesto: no se quedan sin pérdida auxiliar del todo. Suman una muy chica, con alfa cero coma cero cero cero uno, solo para evitar desbalances extremos dentro de una misma secuencia.

Resumiendo

• L_aux suma su gradiente al de la tarea: un compromiso con α
• Sesgo b_i: entra en el top-k, no en el peso g
• Regla: b_i ← b_i + u·sign(c̄ − c_i), sin gradiente
• MaxVio = (máx c − c̄) / c̄
• Juguete: MaxVio 0,37 con L_aux, 0,09 con sesgo
• Próximo: abrir Granite y ver qué aprende cada experto

Papers del episodio

• Wang et al., 2024 — Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts
• DeepSeek-AI, 2024 — DeepSeek-V3 Technical Report