Estabilidad

Optimizadores, con foco en AdamW. Episodio nueve: estabilidad. Ya tenemos el optimizador y el schedule. Hoy, tres detalles que separan un entrenamiento que termina de uno que explota o se congela: gradient clipping, el épsilon de Adam, adaptive moment estimation, y entrenar con números de dieciséis bits.

Tres formas de romper un entrenamiento

• Un gradiente gigante, de un batch roto
• Dividir por casi cero
• Números que no entran en 16 bits
• Las tres se ven igual: inf, NaN o una pérdida congelada

Primero, un batch raro puede dar un gradiente mil veces más grande que el normal. Segundo, Adam divide por la raíz de v: si v es casi cero, el cociente se dispara. Para eso está épsilon. Tercero, para ahorrar memoria y tiempo se entrena con números de dieciséis bits, que tienen menos rango y menos precisión que los de treinta y dos. Las tres terminan parecido: infinito, NaN, que quiere decir not a number, o una pérdida que deja de bajar. Vamos una por una.

Parte 1 · Gradient clipping

Si el gradiente es demasiado largo, se lo acorta. La dirección no se toca.

clip_grad_norm_, a mano y en torch

En Álgebra lineal para LLMs, episodio diez, usamos esta función. Ahora, por dentro. Dos parámetros: a, con gradiente tres y cuatro, y b, con doce. La norma se calcula sobre todos los parámetros juntos, como un solo vector largo. El factor es el máximo permitido, uno, dividido la norma; si la norma ya era chica, el factor queda en uno y no se toca nada. Y la función de PyTorch, con max norm uno. La norma total es trece. El factor, cero coma cero siete siete. A mano y en torch, el mismo vector: cero coma dos tres, cero coma tres uno, cero coma nueve dos. Y clip grad norm devuelve trece: la norma de antes de recortar, que conviene registrar en cada paso.

a = torch.zeros(2, requires_grad=True); a.grad = torch.tensor([3.0, 4.0])
b = torch.zeros(1, requires_grad=True); b.grad = torch.tensor([12.0])

# a mano: una sola norma para todos los parámetros juntos
total = torch.cat([a.grad, b.grad]).norm()           # √(3² + 4² + 12²)
factor = min(1.0, 1.0 / (total.item() + 1e-6))
print("norma total:", total.item(), " factor:", round(factor, 4))
print("a mano:", [round(x * factor, 4) for x in [3.0, 4.0, 12.0]])

devuelto = torch.nn.utils.clip_grad_norm_([a, b], max_norm=1.0)
print("torch :", [round(x, 4) for x in torch.cat([a.grad, b.grad]).tolist()])
print("clip_grad_norm_ devolvió:", devuelto.item())

norma total: 13.0  factor: 0.0769
a mano: [0.2308, 0.3077, 0.9231]
torch : [0.2308, 0.3077, 0.9231]
clip_grad_norm_ devolvió: 13.0

El recorte por norma, con 3, 4 y 12

$\|g\|$: la norma de todo el gradiente, todos los parámetros juntos

$c$: max_norm, acá 1

cambia el largo, no la dirección

La norma, con las dobles barras, es el largo del vector: la raíz de la suma de los cuadrados. Juntamos los tres números de a y b. En el dibujo, nueve más dieciséis más ciento cuarenta y cuatro da ciento sesenta y nueve; la raíz, trece. c es max norm, el largo máximo que permitimos, acá uno. Multiplicamos todo el gradiente por el mínimo entre uno y c sobre la norma. Acá es un treceavo. En el dibujo, las tres componentes divididas por trece. Dividir todas las componentes por el mismo número no cambia hacia dónde apunta el vector: solo su largo. En el dibujo, de trece a uno.

Por valor o por norma

PyTorch tiene otra función de nombre parecido, y conviene no confundirlas. Un gradiente de cuatro y dos. clip grad value recorta cada número por separado al intervalo de menos dos a dos. clip grad norm escala el vector entero para que su norma sea dos. Comparamos cada resultado con el original usando el coseno: uno quiere decir misma dirección. Por valor queda dos y dos, coseno cero coma nueve cuatro nueve: la dirección cambió. Por norma queda uno coma siete nueve y cero coma ocho nueve, coseno uno: misma dirección, más corto.

g = torch.tensor([4.0, 2.0])                         # un gradiente grande
p1 = torch.zeros(2, requires_grad=True); p1.grad = g.clone()
p2 = torch.zeros(2, requires_grad=True); p2.grad = g.clone()

torch.nn.utils.clip_grad_value_([p1], clip_value=2.0)  # cada número a [−2, 2]
torch.nn.utils.clip_grad_norm_([p2], max_norm=2.0)     # todo el vector, escalado

for nombre, p in [("por valor", p1), ("por norma", p2)]:
    cos = F.cosine_similarity(p.grad, g, dim=0).item()
    print(f"{nombre}: {[round(x, 3) for x in p.grad.tolist()]}"
          f"  coseno con g = {cos:.3f}")

por valor: [2.0, 2.0]  coseno con g = 0.949
por norma: [1.789, 0.894]  coseno con g = 1.000

Los tres vectores

Dibujados. El gradiente original, cuatro y dos. Recortado por valor: la componente grande se achata y la chica queda igual, así que el vector gira. En el dibujo, el ángulo entre los dos. Ya no es la dirección de máximo descenso. Recortado por norma: sobre la misma recta que el original, solo más corto. Por eso lo habitual al entrenar LLMs es recortar por norma, muchas veces con max norm uno.

Pascanu, Mikolov y Bengio, 2012 — On the difficulty of training Recurrent Neural Networks · arXiv:1211.5063

gradient norm clipping · «We propose a gradient norm clipping strategy to deal with exploding gradients»

El recorte por norma viene de Pascanu, Mikolov y Bengio, en dos mil doce, estudiando redes recurrentes, donde el gradiente puede explotar. Proponen una estrategia de gradient norm clipping para manejar los exploding gradients. No confundir con el recorte de PPO, proximal policy optimization, de Alineación, episodio cuatro: ese recorta un ratio de probabilidades en la pérdida, no el gradiente.

Un batch roto, con y sin clip

¿Hace falta con Adam, que ya normaliza? Veinte pesos que arrancan en tres; la pérdida es la norma al cuadrado, con el mínimo en cero. En el paso cincuenta, un batch roto: el gradiente se multiplica por mil. Una corrida sin recorte, con max norm infinito, y otra con max norm uno. Y leemos la raíz de v sombrero del primer peso. La norma del batch roto es veintidós mil cuatrocientos cincuenta y seis. Hasta el paso cuarenta y nueve, las dos corridas van casi iguales. En el cincuenta y uno, la raíz de v saltó de cinco coma cinco a setecientos cinco. Sin clip, la pérdida queda en ciento quince en el paso trescientos. Con clip, cero coma cero uno.

def entrenar(clip, pasos=300):
    torch.manual_seed(0)
    w = torch.full((20,), 3.0, requires_grad=True)       # la pérdida: ‖w‖²
    opt = torch.optim.Adam([w], lr=0.01)
    hist = []
    for t in range(pasos):
        opt.zero_grad()
        perdida = (w ** 2).sum()
        perdida.backward()
        if t == 50: w.grad *= 1000                        # un batch roto
        norma = torch.nn.utils.clip_grad_norm_([w], clip) # devuelve la norma de antes
        opt.step()
        est = opt.state[w]                                # √v̂ del primer peso
        rv = (est["exp_avg_sq"][0] / (1 - 0.999 ** est["step"])).sqrt()
        hist.append((perdida.item(), norma.item(), rv.item()))
    return hist
sin, con = entrenar(float("inf")), entrenar(1.0)
print("norma en el paso 50:", round(sin[50][1]))
print("paso   pérdida sin clip   con clip   √v̂ sin clip")
for t in (49, 51, 100, 299):
    print(f"{t:4}   {sin[t][0]:16.2f}   {con[t][0]:8.2f}   {sin[t][2]:11.1f}")

norma en el paso 50: 22456
paso   pérdida sin clip   con clip   √v̂ sin clip
  49             127.02     126.00           5.5
  51             125.35     124.00         704.9
 100             118.14      80.00         499.6
 299             115.42       0.01         275.4

La pérdida, antes y después del batch roto

Sin clip: hasta el paso cincuenta baja como la otra. Después del batch roto, casi se aplana: Adam no explotó, se congeló. Con clip: el batch roto entra recortado a norma uno, como todos, y la pérdida sigue bajando hasta casi cero. El daño quedó guardado en v.

Por qué Adam se congela

un solo paso con el gradiente ×1000

el gradiente vuelve a lo normal; la división no

v tarda ~700 pasos en olvidar la mitad

En el paso del pico, v suma cero coma cero cero uno por el gradiente al cuadrado, y el gradiente es mil veces el normal: un millón de veces su cuadrado. En el dibujo, la raíz de v sombrero pasa de cinco coma cinco a setecientos cinco, ciento veintiocho veces más. Al paso siguiente, el gradiente ya es normal, pero Adam lo divide por esa raíz inflada: el paso es más de cien veces más chico. En el dibujo, en el paso cien vale casi quinientos; en el doscientos noventa y nueve, doscientos setenta y cinco. v se multiplica por cero coma nueve nueve nueve en cada paso. Para olvidar la mitad: logaritmo natural de dos sobre menos logaritmo de cero coma nueve nueve nueve, unos seiscientos noventa y tres pasos. En el dibujo, la mitad, después la cuarta parte.

Parte 2 · Épsilon

El número más chico del optimizador, y el primero que desaparece en dieciséis bits.

Qué hace ε en el primer paso

Eta uno, para leer el paso en unidades de eta, y épsilon diez a la menos ocho, el default. Cuatro gradientes, de uno a diez a la menos diez. En el primer paso, m sombrero es g y la raíz de v sombrero es el valor absoluto de g, como vimos en el episodio siete; épsilon se suma abajo. Y guardamos épsilon en dos tipos de número. Con g igual a uno o a diez a la menos cuatro, el paso es eta entero: épsilon no se nota. Con g igual a épsilon, medio paso. Con diez a la menos diez, una centésima. Y diez a la menos ocho en float32 es diez a la menos ocho; en float16, cero.

eta, eps = 1.0, 1e-8
for g in [1.0, 1e-4, 1e-8, 1e-10]:                   # primer paso: m̂ = g, √v̂ = |g|
    paso = eta * g / (abs(g) + eps)
    print(f"g = {g:7.0e}   paso = {paso:.4f} · η")
print(f"1e-8 en float32: {torch.tensor(1e-8, dtype=torch.float32).item():.3g}")
print(f"1e-8 en float16: {torch.tensor(1e-8, dtype=torch.float16).item():.3g}")

g =   1e+00   paso = 1.0000 · η
g =   1e-04   paso = 0.9999 · η
g =   1e-08   paso = 0.5000 · η
g =   1e-10   paso = 0.0099 · η
1e-8 en float32: 1e-08
1e-8 en float16: 0

Épsilon separa dos regímenes

$\epsilon = 10^{-8}$; primer paso

$\gg$: mucho mayor; abajo de ε, Adam se vuelve SGD

en float16, 1e-8 se redondea a 0

El paso, con épsilon abajo. En el dibujo, los cuatro números de la salida. Las flechas dobles son mucho mayor y mucho menor. Si g es mucho mayor que épsilon, el paso es eta por el signo: Adam puro. Si es mucho menor, el paso es eta sobre épsilon por g: un SGD, stochastic gradient descent, con un learning rate cien millones de veces eta. En el dibujo, la transición. El mínimo normal de float16 es seis coma uno por diez a la menos cinco; debajo hay subnormales, menos precisos, hasta unos seis por diez a la menos ocho. Diez a la menos ocho queda debajo de todo: cero. Y sin épsilon, un v nulo divide por cero. Por eso m y v se guardan en float32, aunque el modelo use dieciséis bits.

Parte 3 · Precisión mixta

float16 y bfloat16 usan los mismos 16 bits. Los reparten distinto.

torch.finfo: rango y precisión

torch punto finfo describe cada tipo: bits, máximo, mínimo normal y el épsilon de máquina, la distancia entre uno y el siguiente número representable. Es otro épsilon, no el de Adam. Tres números de prueba: setenta mil, uno coma cero cero uno y diez a la menos seis. Los pasamos a float16. Y a bfloat16. float32: máximo tres coma cuatro por diez a la treinta y ocho. float16: máximo sesenta y cinco mil quinientos cuatro, épsilon una milésima. bfloat16: casi el máximo de float32, pero épsilon ocho milésimas. Setenta mil en float16 es infinito; en bfloat16, setenta mil ciento cuarenta y cuatro. Y uno coma cero cero uno en bfloat16 es uno. Diez a la menos seis sobrevive en float16 con un uno por ciento de error.

for dt in [torch.float32, torch.float16, torch.bfloat16]:
    fi = torch.finfo(dt)
    print(f"{str(dt)[6:]:8}  bits {fi.bits}  máx {fi.max:.6g}  eps {fi.eps:.3g}"
          f"  mín {fi.tiny:.3g}")
x = torch.tensor([70000.0, 1.001, 1e-6])
print("float16 :", [f"{v:.6g}" for v in x.to(torch.float16).tolist()])
print("bfloat16:", [f"{v:.6g}" for v in x.to(torch.bfloat16).tolist()])

float32   bits 32  máx 3.40282e+38  eps 1.19e-07  mín 1.18e-38
float16   bits 16  máx 65504  eps 0.000977  mín 6.1e-05
bfloat16  bits 16  máx 3.38953e+38  eps 0.00781  mín 1.18e-38
float16 : ['inf', '1.00098', '1.01328e-06']
bfloat16: ['70144', '1', '9.98378e-07']

Los mismos 16 bits, repartidos distinto

$s$: signo; $e$: exponente (rango); $f$: mantisa (precisión)

5 bits de exponente contra 8

10 bits de mantisa contra 7

Un número de punto flotante se guarda como un signo, s; un exponente, e, que dice cuántas veces multiplicar por dos; y una mantisa, f, los dígitos detrás de la coma. En el dibujo, bits de exponente: ocho, cinco y ocho. Más bits de exponente, más rango. float16 llega a sesenta y cinco mil quinientos cuatro; bfloat16, casi al mismo que float32. En el dibujo, en potencias de diez: cuatro coma ocho contra treinta y ocho coma cinco. Más bits de mantisa, más precisión. El épsilon de máquina es dos a la menos cantidad de bits de mantisa: una milésima para float16, ocho milésimas para bfloat16. En el dibujo, veintitrés, diez y siete bits. bfloat16 cambia precisión por rango.

autocast y GradScaler

Así se entrena en float16. Corre en CPU; en una GPU, graphics processing unit, es igual. GradScaler multiplica la pérdida por un número grande antes del backward, para que los gradientes chicos no se redondeen a cero en float16. autocast hace el forward en float16 donde conviene; los pesos siguen en float32. El backward con la pérdida escalada. unscale divide los gradientes por la escala: tiene que ir antes del clip, si no recortás números sesenta y cinco mil veces más grandes. Y step saltea el paso si apareció un infinito. La salida es float16 y la escala inicial, sesenta y cinco mil quinientos treinta y seis. La norma escalada, sesenta y cinco mil cuatrocientos veintiséis; la real, cero coma nueve nueve ocho.

red = torch.nn.Linear(16, 1)
opt = torch.optim.AdamW(red.parameters(), lr=1e-3)
scaler = torch.amp.GradScaler("cpu")                  # escala para float16
with torch.autocast("cpu", dtype=torch.float16):     # el forward en float16
    salida = red(X)
    perdida = ((salida.float() - Y) ** 2).mean()
print("salida:", salida.dtype); print("escala:", scaler.get_scale())
scaler.scale(perdida).backward()                      # gradientes × 65536
print("norma escalada:  ", round(red.weight.grad.norm().item(), 1))
scaler.unscale_(opt)                                  # ÷ 65536, antes del clip
print("norma real:      ", round(red.weight.grad.norm().item(), 4))
torch.nn.utils.clip_grad_norm_(red.parameters(), 1.0)
scaler.step(opt)                                      # saltea si hubo inf
scaler.update()

salida: torch.float16
escala: 65536.0
norma escalada:   65426.4
norma real:       0.9983

Micikevicius et al., 2017 — Mixed Precision Training · arXiv:1710.03740

loss scaling · «weights, activations and gradients are stored in IEEE half-precision format» … «we propose scaling the loss appropriately to handle the loss of information with half-precision gradients»

La receta es de Micikevicius y coautores, en dos mil diecisiete. Guardan pesos, activaciones y gradientes en media precisión, que es float16, mantienen una copia de los pesos en precisión simple, y proponen escalar la pérdida para no perder información en los gradientes de media precisión. Dicen que así la memoria baja casi a la mitad.

float16 o bfloat16

Para elegir. float16 tiene más precisión, pero un rango corto: hay que escalar la pérdida y vigilar los infinitos. bfloat16 tiene el rango de float32, así que en general no hace falta escalar, pero redondea más fuerte: uno coma cero cero uno es uno. Por eso, en los dos casos, el optimizador guarda m, v y una copia de los pesos en float32: las actualizaciones chicas no sobrevivirían en dieciséis bits.

Resumiendo

• clip_grad_norm_: escala por $\min(1, c/\|g\|)$; devuelve la norma
• Por norma conserva la dirección; por valor, no
• Un batch ×1000 infla $\sqrt{\hat v}$ 128 veces y congela a Adam
• ε = 1e-8: abajo de ε Adam es SGD; en float16, ε = 0
• float16: más precisión, poco rango; bfloat16: al revés
• Próximo: después de AdamW, Adafactor, Lion y Muon

Papers del episodio

• Pascanu, Mikolov y Bengio, 2012 — On the difficulty of training Recurrent Neural Networks
• Micikevicius et al., 2017 — Mixed Precision Training