Del transformer de 2017 a Gemma

Attention is all you need. Episodio nueve: del transformer de dos mil diecisiete a Gemma. Siete años después, un modelo como Gemma 2 sigue siendo ese transformer, con cinco cambios. Hoy, cada cambio, contra el paper, mirado en la configuración y en los módulos de Gemma 2 2B, con sus pesos reales.

2017 contra Gemma 2

Los cambios, de un vistazo. El paper: encoder y decoder, LayerNorm después de la suma, senos sumados una vez abajo, y una red densa con máximo entre cero y z. Gemma: solo decoder. RMSNorm, root mean square norm, una normalización más barata, antes y después de cada subcapa. RoPE, rotary position embedding: la posición entra rotando la query y la key. GQA, grouped-query attention: ocho heads de query comparten cuatro de key y value. Y una red densa con compuerta.

La configuración de Gemma 2 2B

La configuración, leída del modelo descargado. Las heads de key y value, aparte. La base de RoPE. Y el tamaño del KV cache, que explicamos en la parte cuatro. Veintiséis capas, d sub model dos mil trescientos cuatro, ocho heads de query y cuatro de key y value, de doscientos cincuenta y seis números cada una. La FFN, feed-forward network, de nueve mil doscientos dieciséis, con una GELU, gaussian error linear unit, aproximada. Theta de RoPE, diez mil: la misma base de los senos del paper. Y dos detalles nuevos: la mitad de las capas mira solo los últimos cuatro mil noventa y seis tokens, y los scores se recortan suave en cincuenta.

c = AutoConfig.from_pretrained("google/gemma-2-2b")
print(f"capas {c.num_hidden_layers} · d_model {c.hidden_size} · heads "
      f"{c.num_attention_heads} (de key/value {c.num_key_value_heads})")
print(f"head_dim {c.head_dim} · FFN {c.intermediate_size} · "
      f"activación {c.hidden_activation}")
print(f"vocabulario {c.vocab_size} · RoPE θ = "
      f"{c.rope_parameters['rope_theta']:.0f} · contexto {c.max_position_embeddings}")
print(f"ventana local {c.sliding_window} · soft-cap {c.attn_logit_softcapping}"
      f" y {c.final_logit_softcapping}")

kv = lambda h_kv: 2 * c.num_hidden_layers * h_kv * c.head_dim * 2  # K y V, bf16
print(f"KV cache por token: {kv(4) / 1024:.0f} KiB con 4 heads de K/V,"
      f" {kv(8) / 1024:.0f} KiB con 8")

capas 26 · d_model 2304 · heads 8 (de key/value 4)
head_dim 256 · FFN 9216 · activación gelu_pytorch_tanh
vocabulario 256000 · RoPE θ = 10000 · contexto 8192
ventana local 4096 · soft-cap 50.0 y 30.0
KV cache por token: 104 KiB con 4 heads de K/V, 208 KiB con 8

Parte 1 · Solo decoder

El paper traducía: una frase entra al encoder y otra sale del decoder. Un modelo de lenguaje solo continúa texto.

Qué se va con el encoder

• Ni encoder ni cross-attention: una sola pila
• Cada token mira solo hacia atrás: la máscara causal del ep. 5
• Una tarea: predecir el token siguiente
• Gemma: 26 capas iguales, todas de decoder

Qué cambia al sacar el encoder. No hay memoria que mirar, así que la cross-attention desaparece: cada capa tiene self-attention y FFN, nada más. La self-attention lleva la máscara causal del episodio cinco: cada token ve solo a los anteriores. Y el entrenamiento es uno solo: en cada posición, predecir el token que sigue. Gemma son veintiséis capas iguales de esas. El reporte técnico de Gemma 2 lo dice en una línea: un transformer solo decoder, y cita a Vaswani.

Una capa de Gemma, por dentro

El device meta crea los tensores sin memoria: solo formas. Una capa de decoder de Gemma, la cero. Sus capas lineales. Y sus normalizaciones. q proj lleva de dos mil trescientos cuatro a dos mil cuarenta y ocho: ocho heads por doscientos cincuenta y seis. k proj y v proj, a mil veinticuatro: solo cuatro heads. o proj vuelve. Tres matrices de MLP, multilayer perceptron, con nueve mil doscientos dieciséis neuronas. Y cuatro RMSNorm, no dos: antes y después de la atención, y antes y después del MLP.

with torch.device("meta"):                       # solo formas, sin memoria
    capa = Gemma2DecoderLayer(cfg, layer_idx=0)

for nombre, m in capa.named_modules():
    if isinstance(m, torch.nn.Linear):
        print(f"{nombre:16} {tuple(m.weight.shape)}")
normas = [n.replace("_layernorm", "") for n, _ in capa.named_children()
          if "norm" in n]
print(len(normas), "×", type(capa.input_layernorm).__name__,
      tuple(capa.input_layernorm.weight.shape), ":")
print(" ", *normas)

self_attn.q_proj (2048, 2304)
self_attn.k_proj (1024, 2304)
self_attn.v_proj (1024, 2304)
self_attn.o_proj (2304, 2048)
mlp.gate_proj    (9216, 2304)
mlp.up_proj      (9216, 2304)
mlp.down_proj    (2304, 9216)
4 × Gemma2RMSNorm (2304,) :
  input post_attention pre_feedforward post_feedforward

Parte 2 · RMSNorm, antes y después

LayerNorm resta la media y divide por el desvío. RMSNorm se saltea la resta.

RMSNorm, con los mismos cuatro números

la raíz de la media de los cuadrados

sin restar la media, sin $\beta$; Gemma escala por $1 + w$

con media cero, RMSNorm y LayerNorm son la misma cuenta

El token del episodio seis: dos, cuatro, seis y ocho. RMS es root mean square: la raíz de la media de los cuadrados. Cuatro, dieciséis, treinta y seis y sesenta y cuatro: promedio treinta, raíz cinco coma cuarenta y ocho. Dividimos por eso, y multiplicamos elemento a elemento por uno más w, un vector que se aprende. Sin restar la media y sin beta: una pasada menos sobre los datos. Que Gemma escale por uno más w, y no por w, se ve en su código: w arranca en cero. Y si el vector ya tiene media cero, la raíz de la media de los cuadrados es justo el desvío, y RMSNorm y LayerNorm dan lo mismo.

RMSNorm a mano y con Gemma2RMSNorm

La cuenta a mano. El módulo de Gemma, recién creado. Con w igual a uno, para ver el uno más w. Y un vector centrado, contra nn punto LayerNorm. RMS, raíz de treinta: cinco coma cuatro siete siete. El vector normalizado, igual al del módulo: True. Con w igual a uno, todo se duplica: el factor es uno más uno. Y con media cero, RMSNorm igual a LayerNorm: True.

x = torch.tensor([2., 4., 6., 8.])                 # el token del episodio 6
rms = torch.sqrt((x ** 2).mean() + 1e-6)          # raíz de la media de x²
print(f"RMS = √30 = {rms:.3f} → x / RMS =", [round(v, 3) for v in (x / rms).tolist()])

norma = Gemma2RMSNorm(4)                           # peso w = 0 al inicio
print("igual a Gemma2RMSNorm:", torch.allclose(norma(x), x / rms))
with torch.no_grad():
    norma.weight.fill_(1.0)                        # Gemma escala por (1 + w)
print("con w = 1:", [round(v, 3) for v in norma(x).tolist()])

c = x - x.mean()                                   # un vector con media 0
print("media 0 → RMSNorm == LayerNorm:",
      torch.allclose(Gemma2RMSNorm(4, eps=1e-5)(c), nn.LayerNorm(4)(c), atol=1e-5))

RMS = √30 = 5.477 → x / RMS = [0.365, 0.73, 1.095, 1.461]
igual a Gemma2RMSNorm: True
con w = 1: [0.73, 1.461, 2.191, 2.921]
media 0 → RMSNorm == LayerNorm: True

Una capa de Gemma 2

La capa entera. Entra x, el residual stream. Se normaliza la entrada, pasa por la atención, se normaliza la salida, y recién ahí se suma. Lo mismo con el MLP. Comparalo con el episodio seis: el paper normalizaba después de la suma, y el residual stream pasaba por cada LayerNorm. Acá la suma queda limpia, como en pre-norm, y además cada subcapa entrega algo de escala controlada. El reporte de Gemma 2 dice que lo hace para estabilizar el entrenamiento.

x → RMSNorm → atención → RMSNorm → x + · → RMSNorm → MLP → RMSNorm → x + ·

Parte 3 · RoPE: rotar en vez de sumar

En el episodio cuatro vimos que mover una posición es rotar un par de senos. RoPE usa esa rotación directamente sobre la query y la key.

RoPE, de un par a todos

un par de números, girado $m\theta$: $m$ es la posición

$R_m^\top R_n$ es girar $(n - m)\theta$: solo importa la distancia

cada par gira a su propia velocidad

$d/2$ rotaciones de $2 \times 2$ en la diagonal

El paper de RoFormer arranca con dos dimensiones. f de x y m: girar un par de números un ángulo m por theta, donde m es la posición del token y theta una frecuencia fija. El score entre una query en la posición m y una key en la n: la transpuesta de una rotación es la rotación al revés, así que R m transpuesta por R n es girar n menos m. El producto punto depende solo de la distancia. Con d números se arman d sobre dos pares, cada uno con su theta: la misma escala geométrica de los senos del episodio cuatro. Con d igual a ocho, uno, cero coma uno, cero coma cero uno y cero coma cero cero uno. Todo junto: bloques de dos por dos en la diagonal. Se rotan q y k en cada capa; nada se suma al embedding.

Girar la query según su posición

La rotación, en el plano, con theta igual a uno. En verde, la query. En la posición uno, todo el plano gira un radián, en sentido contrario a las agujas del reloj: la grilla no se deforma y q no cambia de largo. En la posición dos, gira dos radianes. Si la key también gira según su posición, lo que queda entre las dos es solo la diferencia de ángulos.

RoPE en doce líneas

Una theta por par, con base diez mil. El ángulo de cada par: posición por theta. La rotación de dos por dos, escrita con las columnas pares e impares. Y la prueba: el score entre una query y una key en cuatro pares de posiciones. Tres, uno; diez, ocho; y cien, noventa y ocho: distancia dos en los tres casos, y el mismo score, menos diez coma uno cuatro dos siete. Con distancia cuatro, otro valor. Y la norma no cambia al rotar: ocho coma tres siete cero cinco antes y después.

def rope(x, pos, base=10000.0):                  # x: [T, d], d par
    d = x.shape[-1]
    theta = base ** (-torch.arange(0, d, 2) / d)  # una frecuencia por par
    ang = pos[:, None] * theta                    # [T, d/2]: pos · θ_i
    x1, x2 = x[..., 0::2], x[..., 1::2]           # pares (0,1), (2,3), …
    out = torch.empty_like(x)
    out[..., 0::2] = x1 * ang.cos() - x2 * ang.sin()
    out[..., 1::2] = x1 * ang.sin() + x2 * ang.cos()
    return out

q, k = torch.randn(64), torch.randn(64)          # una query y una key
def score(m, n):                                 # query en m, key en n
    return rope(q[None], torch.tensor([m]))[0] @ rope(k[None], torch.tensor([n]))[0]

for m, n in [(3, 1), (10, 8), (100, 98), (5, 1)]:
    print(f"m = {m:3}, n = {n:3}, m − n = {m - n}: q·k = {score(m, n):.4f}")
print("norma antes y después:", f"{q.norm():.4f}",
      f"{rope(q[None], torch.tensor([50]))[0].norm():.4f}")

m =   3, n =   1, m − n = 2: q·k = -10.1427
m =  10, n =   8, m − n = 2: q·k = -10.1427
m = 100, n =  98, m − n = 2: q·k = -10.1427
m =   5, n =   1, m − n = 4: q·k = -7.5359
norma antes y después: 8.3705 8.3705

El RoPE de Gemma es el del paper, reordenado

Ahora contra el código de Gemma. El módulo de RoPE de Gemma, con su configuración. La función que aplica la rotación. Ojo: no rota pares vecinos, cero con uno, dos con tres, sino la primera mitad contra la segunda, el número i con el i más ciento veintiocho. Una permutación lleva de una convención a la otra. Y aplicamos nuestra versión del paper sobre la query reordenada. Theta diez mil; cos y sin de seis posiciones por doscientos cincuenta y seis. Y Gemma igual al RoPE del paper, reordenado: True. Si alguna vez portás pesos entre implementaciones, esta permutación es la trampa clásica.

emb = Gemma2RotaryEmbedding(cfg)                     # el de Gemma 2 2B
T, d = 6, cfg.head_dim                               # head_dim = 256
q, k = torch.randn(1, 1, T, d), torch.randn(1, 1, T, d)
pos = torch.arange(T)
cos, sin = emb(q, pos[None])
q_hf, k_hf = apply_rotary_pos_emb(q, k, cos, sin)    # mitades: (i, i + d/2)

perm = torch.cat([torch.arange(0, d, 2), torch.arange(1, d, 2)])
q_paper = rope(q[0, 0, :, perm.argsort()], pos)[:, perm]   # pares vecinos
print("theta de Gemma:", cfg.rope_parameters["rope_theta"],
      "· cos, sin:", tuple(cos.shape))
print("Gemma == RoPE del paper, reordenado:",
      torch.allclose(q_hf[0, 0], q_paper, atol=1e-5))

theta de Gemma: 10000.0 · cos, sin: (1, 6, 256)
Gemma == RoPE del paper, reordenado: True

Su et al., 2021 — RoFormer: Enhanced Transformer with Rotary Position Embedding · arXiv:2104.09864

θᵢ = 10000^(−2(i−1)/d) · «RoPE encodes the absolute position with a rotation matrix and meanwhile incorporates the explicit relative position dependency in self-attention formulation.»

Así lo resumen sus autores. RoPE codifica la posición absoluta con una matriz de rotación, y al mismo tiempo hace que la atención dependa explícitamente de la posición relativa. Las dos cosas que vimos: cada token gira según dónde está, y el score solo ve la distancia.

Parte 4 · GQA y la FFN con compuerta

Los dos últimos cambios: compartir keys y values entre heads, y una FFN con una compuerta.

GQA: menos keys, menos memoria

4 heads de key, cada una para 2 heads de query

K y V de cada token, en las $L = 26$ capas, en bf16

GQA-$G$: $G$ grupos; Gemma tiene $G = 4$

GQA, que adelantamos en el episodio tres. Gemma calcula cuatro heads de key y de value, y cada una la usan dos heads de query: repeat interleave las duplica. ¿Para qué? Al generar, el modelo guarda las keys y values de todos los tokens anteriores, el KV cache, key-value cache, para no recalcularlas. Su tamaño es dos, por key y value, por L capas, por h sub kv heads, por d sub head números, por dos bytes cada uno: ciento cuatro kibibytes por token, la mitad que con ocho. El paper de GQA llama G a los grupos: G igual a uno es multi-query, una sola key; G igual a h es la multi-head original. Gemma tiene cuatro grupos, aunque su reporte escribe num groups igual a dos: parece contar las queries de cada grupo.

Ainslie et al., 2023 — GQA: Training Generalized Multi-Query Transformer Models · arXiv:2305.13245

GQA-G: G grupos · «Grouped-query attention divides query heads into G groups, each of which shares a single key head and value head.»

La definición del paper. Grouped-query attention divide las heads de query en G grupos, y cada grupo comparte una sola head de key y una de value. Lo presentan como un punto medio: casi la calidad de multi-head, casi la velocidad de una sola key.

La FFN con compuerta

GeGLU: dos proyecciones hacia arriba, una hace de compuerta

la compuerta decide cuánto de $u$ pasa

por capa; el reporte de Gemma 2 dice «18432» = $2 \cdot 9216$

El MLP de Gemma, que en Alineación, episodio uno, vimos de pasada. Dos proyecciones hacia arriba: W gate y W up. La de gate pasa por una GELU, una no linealidad suave parecida a la ReLU, rectified linear unit. En el dibujo, una neurona de la compuerta: z vale uno, y la GELU da cero coma ochenta y cuatro. Al conjunto le dicen GeGLU. Las dos se multiplican elemento a elemento: la compuerta decide cuánto de u pasa. Con cero coma ochenta y cuatro y u igual a dos, pasa uno coma sesenta y ocho; con casi cero, se cierra. Son tres matrices en vez de dos: sesenta y tres coma siete millones de parámetros por capa, contra dos coma uno del paper. El reporte de Gemma 2 escribe dieciocho mil cuatrocientos treinta y dos: es gate más up, dos por nueve mil doscientos dieciséis.

El MLP de la capa 0, con sus pesos reales

Cargamos las tres matrices de la capa cero de Gemma 2 2B desde los archivos del modelo. La compuerta: GELU con la aproximación tanh, la que dice la configuración. Por u, y W down. Y contamos. La cuenta a mano coincide con Gemma2MLP: True. Tres tokens de dos mil trescientos cuatro, compuerta y u de nueve mil doscientos dieciséis, y de vuelta. Sesenta y tres coma siete millones de parámetros, y dos por nueve mil doscientos dieciséis es dieciocho mil cuatrocientos treinta y dos: el número de la tabla del reporte.

mlp = Gemma2MLP(cfg)                                     # pesos reales, capa 0
for p in ("gate_proj", "up_proj", "down_proj"):
    getattr(mlp, p).weight.data = peso(f"model.layers.0.mlp.{p}.weight")

x = torch.randn(3, cfg.hidden_size)                      # 3 tokens
g = F.gelu(x @ mlp.gate_proj.weight.T, approximate="tanh")   # la compuerta
u = x @ mlp.up_proj.weight.T                             # lo que pasa
a_mano = (g * u) @ mlp.down_proj.weight.T
print("a mano == Gemma2MLP:", torch.allclose(a_mano, mlp(x), atol=1e-4))
print("x", tuple(x.shape), "→ g y u", tuple(g.shape), "→ salida", tuple(a_mano.shape))

n = sum(p.numel() for p in mlp.parameters())
print(f"parámetros: {n / 1e6:.1f} M (el paper: 2.1 M) · 2 × 9216 = {2 * 9216}")

a mano == Gemma2MLP: True
x (3, 2304) → g y u (3, 9216) → salida (3, 2304)
parámetros: 63.7 M (el paper: 2.1 M) · 2 × 9216 = 18432

Resumiendo

• Solo decoder: una pila de 26 capas con máscara causal
• RMSNorm: $x / \mathrm{RMS}(x) \odot (1 + w)$, antes y después de cada subcapa
• RoPE: rota q y k; $q_m^\top k_n$ solo depende de $m - n$
• GQA: 8 queries, 4 keys; KV cache de 104 KiB por token
• FFN con compuerta: GELU(gate) ⊙ up, 63,7 M por capa
• Próximo: un transformer entero, desde cero

Papers del episodio

• Su et al., 2021 — RoFormer: Enhanced Transformer with Rotary Position Embedding
• Ainslie et al., 2023 — GQA: Training Generalized Multi-Query Transformer Models