De GPT-2 a Llama: cuatro cambios y el motivo de cada uno

La parte 6 acababa con las cinco decisiones de GPT-2 y el bucle que lo sirve. Si abres un modelo de esta década, el bloque sigue siendo el de la parte 5, con dos etapas que leen cada una su copia normalizada del bus y le suman el resultado. Por dentro, eso sí, han cambiado cuatro cosas.

He leído modeling_llama.py de transformers 5.17 comparándolo con el model.py de nanoGPT y he cargado SmolLM2-135M, un modelo con la arquitectura de Llama y el tamaño de GPT-2 small. Después he ido deshaciendo cada cambio sobre sus pesos entrenados y midiendo qué se rompía.

El mismo bloque, cuatro retoques

Resumido, LlamaDecoderLayer.forward (modeling_llama.py:305-323) queda así:

residual = hidden_states
hidden_states = self.input_layernorm(hidden_states)            # RMSNorm, not LayerNorm
hidden_states, _ = self.self_attn(hidden_states, position_embeddings=...)
hidden_states = residual + hidden_states                       # the bus, as in GPT-2
residual = hidden_states
hidden_states = self.mlp(self.post_attention_layernorm(hidden_states))
hidden_states = residual + hidden_states

Son las mismas cuatro líneas de la parte 5, solo que con más nombres. Las sumas siguen donde estaban, y también el pre-LN, es decir, la norma que se aplica antes de cada etapa y no después. Lo que ha cambiado está dentro de las llamadas: cómo entran las posiciones en la atención, qué calcula la norma, qué forma tiene la red feed-forward y cuántas keys y values guarda la atención. (Llama, además, quita los vectores de sesgo de sus proyecciones lineales, aunque nada de lo que viene después depende de ellos.)

La comparación entre los dos modelos es justa. GPT-2 small tiene 124 millones de parámetros repartidos en 12 bloques de anchura 768, y SmolLM2-135M, 135 millones en 30 bloques de anchura 576. Ninguno de los cuatro cambios tiene que ver con el tamaño, así que podrías construir cualquiera de los dos modelos con ellos o sin ellos.

Las posiciones pasan a ser una rotación

GPT-2 suma a cada token una fila aprendida, distinta para cada posición, y la tabla se acaba en 1024. Llama no tiene tabla de posiciones, y lo que hace es girar la query y la key (modeling_llama.py:158-159).

q_embed = (q * cos) + (rotate_half(q) * sin)
k_embed = (k * cos) + (rotate_half(k) * sin)

Es el embedding posicional rotatorio, RoPE. Como se vio en la parte 4, la atención puntúa un par de tokens comparando un vector query con un vector key. Cada query y cada key se divide en cabezas de 64 números, y el vector de cada cabeza se lee como 32 pares. Cada par es un punto sobre una circunferencia y gira, como la aguja de un reloj, un ángulo proporcional a la posición. rotate_half es la manera que tiene el código de escribir ese giro sin una matriz. Las velocidades son distintas para cada par y todas salen de un solo número de la configuración, la base de rotación, que en SmolLM2 vale 100 000. Con esa base, el par más rápido da una vuelta cada 6,3 posiciones y el más lento, una cada 438 000.

La rotación no está de adorno. Si giras una query según la posición m y una key según la posición n, su producto escalar (la puntuación de la parte 4) depende solo de m − n. He pasado una query y una key aleatorias por la rotación de SmolLM2. La puntuación ha dado −13,4724 en las posiciones 10 y 3, y −13,4723 en 1010 y 1003. Con una posición menos de distancia, en 10 y 4, ha dado −10,786.

El modelo entero hereda esa propiedad. He sumado 500 a todos los ids de posición de un texto de 512 tokens, y después 5000. La pérdida de SmolLM2 sobre ese texto (su sorpresa media ante el token siguiente) ha sido de 3,381 sin desplazar y de 3,381 tras cada desplazamiento. A GPT-2, el mismo desplazamiento de 500 lo ha llevado de 3,85 a 8,97 porque cada token pasa a leer otra fila de su tabla. Al quitar la rotación del todo, con todas las posiciones a 0, SmolLM2 ha pasado de 3,12 a 6,28 en el texto completo de 1024 tokens. Y bajar la base de rotación de los 100 000 de SmolLM2 a 10 000, el valor del artículo original de RoPE y de Llama 2, ha costado un 3 %.

Arrastra la posición de la key para ver cómo cambia la puntuación, y después desplaza las dos para comprobar que no se mueve. Si eliges el par 17, el más lento de los tres, las posiciones cercanas apenas se distinguen. Aquí la puntuación es solo lo que aporta un par, y va de −1 a 1.

Y sin tabla tampoco hay última fila. La configuración de Llama 3.1 8B registra una longitud original de 8192 posiciones y un máximo de 131 072; el modelo llega hasta ahí estirando la rotación. Sus cabezas tienen 128 números, o sea, 64 pares. Los 29 más lentos giran 8 veces más despacio, los 29 más rápidos se quedan igual y los seis intermedios mezclan los dos tratamientos. Meta, además, preentrenó la familia con longitudes crecientes hasta 128K, en seis etapas en el caso del modelo de 405B. No hay ningún assert que salte: pasada la longitud con la que se entrenó un modelo, la calidad hay que medirla, porque nada la garantiza.

La norma deja de centrar

LayerNorm resta a cada fila su media, la divide por su desviación típica y después la escala y la desplaza. LlamaRMSNorm, en modeling_llama.py:62-67, se salta la media y el desplazamiento y divide por la media cuadrática (RMS), que se calcula elevando al cuadrado cada número, promediando y sacando la raíz. Si alguna vez has estandarizado una columna de un dataframe, has restado y has dividido; RMSNorm se queda solo con la división.

variance = hidden_states.pow(2).mean(-1, keepdim=True)
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
return self.weight * hidden_states.to(input_dtype)

RMSNorm apuesta a que la media es pequeña de por sí, y en SmolLM2 lo es. En una capa de norma típica, la media de una fila típica es el 2 % de su media cuadrática, y en la peor capa, el 4,5 %. Si vuelves a poner la LayerNorm completa con la misma escala aprendida, la pérdida sube menos de un 2 %, de 3,124 a 3,175.

Para ser justos con el argumento a su favor, hay que ponerlo a prueba. El artículo de RMSNorm la vende como más barata porque se ahorra el desplazamiento y calcula una estadística donde LayerNorm calcula dos. En la CPU de mi portátil, el F.layer_norm de torch ha tardado entre 50 y 75 microsegundos sobre un tensor [1, 1024, 576], según la ejecución, y F.rms_norm, entre 110 y 135. El profiler dice por qué: en CPU, rms_norm sigue ejecutándose como una cadena de operaciones pequeñas (cuadrado, media, suma, raíz y dos multiplicaciones), mientras que layer_norm es un único kernel optimizado, una sola rutina compilada que hace todo el trabajo.

Con un solo hilo, en cambio, layer_norm ha tardado unos 155 microsegundos y rms_norm, unos 130, así que el orden se invierte. Menos aritmética no se traduce en menos tiempo hasta que alguien escriba el kernel, y en GPU no he medido.

La red feed-forward gana una puerta

La red feed-forward de GPT-2 tiene dos matrices con una GELU en medio, y la de Llama, tres (modeling_llama.py:175).

down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))

La entrada pasa por dos proyecciones en paralelo. Una de ellas atraviesa SiLU, una activación suave parecida a GELU, y después multiplica a la otra elemento a elemento. Esa multiplicación hace de puerta. La combinación se llama SwiGLU y viene de un artículo de 2020 que la probó y publicó mejores resultados sin ofrecer ninguna teoría que los explicara.

Con la misma anchura, tres matrices cuestan más que dos, así que la anchura se reduce. Si d es la anchura del modelo, una red feed-forward convencional de anchura 4d tiene 8d² números sin contar los sesgos. Tres matrices de anchura f tienen 3·d·f, y las dos cuentas coinciden con f = 8d/3. SmolLM2 usa exactamente esa anchura (1536 son 8/3 de 576), y los 11 008 de Llama 2 7B salen de redondear 8/3 de 4096 al múltiplo de 256 superior.

Con 8d/3, la puerta no altera el presupuesto: con cabezas de atención completas la red feed-forward sigue ocupando dos tercios del bloque. El 75 % de SmolLM2 se debe a que comparte keys y values entre cabezas de query, y de eso trata la sección siguiente. Llama 3.1 8B mueve el reparto por los dos lados, con una red feed-forward más ancha, f = 3.5d, y una atención más pequeña porque comparte keys y values. Las dos cosas juntas dan el 81 % que citaba la parte 5 (176 de los 218 millones de parámetros de cada bloque). Solo con la red más ancha y cabezas de atención completas, sale el 72 %.

La puerta importa. Si en SmolLM2 cambias SiLU por GELU, la pérdida pasa de 3,12 a 3,79. Si quitas del todo la multiplicación de la puerta, sube hasta 17,7 y supera con mucho los 10,8 de un modelo que reparte la probabilidad por igual entre los 49 152 tokens.

Es también en la red feed-forward donde entra la mezcla de expertos, con muchas redes feed-forward por bloque y un router, una capa pequeña que elige unas pocas para cada token. Llama 3 siguió siendo denso, a diferencia de Llama 4 y de otras familias.

Se comparten keys y values

GPT-2 tiene tantas cabezas de key y de value como de query. SmolLM2 tiene 9 cabezas de query y 3 cabezas key-value, y cada cabeza key-value sirve a un grupo de tres queries. Es como un índice que atiende tres consultas, con los mismos datos y preguntas distintas. En PyTorch puro, repeat_kv (modeling_llama.py:179-189) copia las tres cabezas de key y value hasta tener nueve antes de hacer las cuentas; el kernel de atención fusionado de torch puede ahorrarse la copia cuando no hay máscara de padding. A esto se le llama grouped-query attention, GQA. Llama 2 solo la usaba en sus dos tamaños más grandes, el de 70B y uno de 34B que no llegó a publicarse, y en Llama 3 la llevan todos los tamaños.

En inferencia, compartir cabezas no es ninguna aproximación. He construido una copia de SmolLM2 con nueve cabezas completas de key y value, en la que cada cabeza repite la de su grupo, y he ejecutado los dos modelos. Los logits han coincidido hasta el último bit, y la copia tiene 147,8 millones de parámetros frente a 134,5. Lo que GQA sacrifica lo sacrifica en el entrenamiento, porque hay menos keys y values distintos que aprender.

Donde se ahorra es en la caché de la parte 6. SmolLM2 guarda 11 520 números por token (la cifra la he sacado de la propia caché de transformers después de 256 tokens), y con cabezas completas guardaría 34 560. GPT-2 small, con 12 bloques frente a los 30 de SmolLM2, guarda 18 432.

Con el tamaño de Llama 3.1 8B, esa proporción decide qué cabe. El modelo tiene 32 cabezas de query y 8 cabezas key-value, así que cada token cuesta 131 KB en bf16, el formato de 16 bits de la parte 3. Con el máximo de 131 072 posiciones, la caché de una sola secuencia ocupa 17,2 GB, frente a los 16,1 GB de los pesos, y con cabezas completas serían 68,7 GB.

Empieza por GPT-2 small y activa los tres cambios de uno en uno. Verás que la parte que ocupa la red feed-forward apenas se mueve. Después baja las cabezas key-value hasta 1 y fíjate en la caché y en esa parte.

Qué vigila el ingeniero de ML

Son cuatro hábitos, y ninguno necesita una GPU.

Recalcula la caché para cada modelo. El número de parámetros no sirve para predecirla, porque la caché son 2 × bloques × cabezas key-value × anchura de cabeza números por token y las cabezas key-value son un campo de la configuración. SmolLM2 es más grande que GPT-2 small y aun así guarda tres octavos menos por token.

Porta la arquitectura antes que los pesos. Cada uno de los cambios anteriores forma parte del contrato con el que se entrenaron los pesos. La activación, el esquema de posiciones con su base de rotación y su escalado, y el número de cabezas key-value son campos de la configuración. El tipo de norma, en cambio, viene con el nombre de la arquitectura. Si te equivocas en el número de cabezas, la carga falla porque las formas no encajan. Si aciertas el número pero te equivocas en la agrupación (repites la serie entera de cabezas compartidas cuando había que repetir cada una varias veces seguidas), no salta ninguna excepción, pero la pérdida ha subido por encima de 11. Lo demás falla con el mismo silencio. Una LayerNorm sin sesgo donde iba RMSNorm ha costado un 1,6 %, GELU donde iba SiLU, un 21 %, y la base de rotación equivocada, un 3 %. Con los pares de rotación formados por números contiguos, cuando el código los forma por mitades, la pérdida ha llegado a 7,24. Si portas un modelo a mano, compara tus logits con los de una implementación de referencia antes de fiarte de él.

Prueba las longitudes a las que sirves el modelo. La rotación quita el límite fijo, pero la pregunta sigue en pie. Si usas un modelo más allá de la longitud con la que se entrenó, te toca a ti medir cómo se comporta ahí, con tus propios datos.

Cuenta aparte los parámetros activos. Con mezcla de expertos, Qwen3-30B-A3B tiene 30 500 millones de parámetros, usa 3350 millones por token y el 95 % de los pesos son expertos. La memoria depende de la primera cifra y la aritmética por token, de la segunda, y pagas las dos.

Esto tampoco es trabajo de modelado. Un bloque de Llama es un bloque de GPT-2 con cuatro piezas cambiadas. Cada cambio es una línea de configuración o de código que puedes leer y cuyo coste puedes calcular, y cualquier desajuste se puede detectar con un test antes de servir la primera petición.