GPT-2 es el transformer más cinco decisiones, y la inferencia es otro programa

La parte 5 terminó con un bloque y una pila de doce. GPT-2 es esa pila con cinco decisiones alrededor, y el mismo checkpoint ejecuta dos programas distintos. El entrenamiento puntúa una secuencia entera en una sola llamada. La generación (lo que al servir el modelo se llama inferencia) escribe los tokens de uno en uno, en un bucle, y de ahí sale la mayoría de los fallos en producción.

He vuelto a cargar los pesos de GPT-2 small de OpenAI en el model.py de nanoGPT y he ejecutado con ellos los dos programas. Después he añadido una caché al bucle, he pasado los pesos a bf16 y he comparado lo que salía.

Cinco decisiones alrededor de la pila

GPTConfig, en model.py:108-116, cabe en una pantalla: 1024 posiciones, 12 bloques, 12 cabezas y anchura 768. El campo del vocabulario vale 50 304 porque, según el comentario de esa línea, se rellena hasta un múltiplo de 64 para ganar velocidad. El checkpoint de OpenAI tiene 50 257 filas, y al cargarlo la configuración se ajusta a él. Alrededor de esos números GPT-2 tomó cinco decisiones, y tres de ellas ya han salido en esta serie.

El tokenizador es el BPE sobre bytes de la parte 1; como codifica cualquier cadena, no existe el token desconocido. La normalización es pre-LN, como en la parte 5. En model.py:138 la tabla de tokens se comparte con la capa de salida. Así, una sola matriz, la tabla que en el recuento de la parte 5 era el 31 % del modelo, sirve para los dos extremos. Con dos matrices separadas, GPT-2 small tendría 163 millones de parámetros, frente a los 124 millones que tiene.

La cuarta decisión son las posiciones, que GPT-2 aprende. wpe es una tabla de forma [1024, 768], con una fila por posición que se suma a la del token antes del primer bloque. Y como toda tabla tiene una última fila, si le pides a forward 1100 posiciones, el assert de model.py:173 te para:

AssertionError: Cannot forward sequence of length 1100, block size is only 1024

La quinta es la inicialización. Cada bloque escribe dos veces en el bus (el nombre que la parte 5 da al flujo residual), con lo que una pila de L bloques acumula 2L escrituras. En nanoGPT todas las matrices, incluidas las tablas de tokens y de posiciones, arrancan con valores aleatorios de desviación típica 0,02 (el código de OpenAI usaba 0,01 para la de posiciones). Los sesgos empiezan en cero, y LayerNorm, en uno. model.py:143-145, en cambio, inicializa con 0.02/√(2L) las dos capas que hacen esas escrituras, para que la suma no crezca con la profundidad. Lo he medido en modelos sin entrenar. Con la inicialización escalada, la norma del flujo se ha multiplicado más o menos por 11 tanto a lo largo de 12 bloques como de 48. Con la inicialización sin escalar, se ha multiplicado por 55 y por 112. A partir de ahí, el entrenamiento hace el resto.

El entrenamiento es una sola llamada

Si le das a forward una secuencia y sus objetivos (la misma secuencia desplazada una posición), predice todas las posiciones a la vez. Le he pasado 92 posiciones de texto y ha devuelto logits de forma [1, 92, 50257], es decir, puntuaciones en bruto para cada token en cada posición, que la softmax convierte en probabilidades. También ha devuelto una sola pérdida, la entropía cruzada de model.py:187, que es menos el logaritmo de la probabilidad asignada al siguiente token correcto, promediado sobre las 92. La máscara de la parte 4 es la que impide hacer trampa, porque la fila t solo ve las filas 0..t y así ninguna predicción puede leer su propia respuesta.

Ese programa se ejecuta en cada paso del entrenamiento, durante cientos de miles de pasos. Como todas las posiciones se conocen de antemano, atraviesan la pila juntas, en unas pocas multiplicaciones de matrices grandes. Si alguna vez has escrito un proceso por lotes sobre una tabla, ya sabes cómo es este programa. En la CPU de mi portátil, una pasada hacia delante sobre 1024 posiciones ha tardado 0,23 segundos.

Sin objetivos, forward toma un atajo en model.py:190 y calcula los logits solo de la última posición. En esa línea asoma por primera vez el segundo programa.

La generación es un bucle

Para escribir texto, el modelo tiene que ver cada token antes de poder predecir el siguiente, y entre posiciones no queda nada que ejecutar en paralelo. Este es, resumido, el generate de nanoGPT en model.py:312-328:

for _ in range(max_new_tokens):
    idx_cond = idx[:, -1024:]                 # crop to the position table
    logits, _ = self(idx_cond)                # the whole sequence, again
    logits = logits[:, -1, :] / temperature   # [B, 50257], last position only
    probs = F.softmax(logits, dim=-1)         # top_k masking skipped here
    idx_next = torch.multinomial(probs, 1)    # draw one token id
    idx = torch.cat((idx, idx_next), dim=1)   # append it and go round

Mira otra vez la línea self(idx_cond). En cada paso, el bucle repite entera la pasada hacia delante sobre la secuencia acumulada hasta entonces y se queda con una sola fila. Si alguna vez has releído un fichero entero para añadirle una línea, ya has escrito este bucle. Con un prompt de 11 tokens, generar 256 tokens nuevos hace pasar 35 456 posiciones por la pila para sacar 256 filas útiles.

La solución es una caché KV. En la parte 4, cada posición recibió un vector key, con el que se comparan las queries posteriores, y un vector value, que es lo que esas queries promedian. Con la máscara de la parte 4, los vectores key y value de la posición t dependen solo de las posiciones 0..t, por lo que, una vez calculados, no vuelven a cambiar. Si los guardas (un par de tensores por bloque), en cada paso basta con pasar el token nuevo por la pila y calcular su atención sobre las filas guardadas. Para los mismos 256 tokens, eso son 266 posiciones. A cambio se paga en memoria, unos 74 kB por posición en GPT-2 small con fp32. Como nanoGPT no trae caché, he escrito una propia, de unas treinta líneas, que reutiliza sus módulos.

El cronómetro lo confirma, aunque no en la misma proporción. Al generar 64, 256 y 512 tokens sin la caché, por la pila han pasado 37, 133 y 261 veces más posiciones, pero el tiempo se ha multiplicado por 2,9, por 5,0 y por 7,4. Con batch 1, cada paso sigue teniendo que leer los 124 millones de pesos para producir un token, y ese coste la caché no lo reduce. Escribir 1023 tokens con la caché ha llevado 9,3 segundos, unas cuarenta veces lo que tarda la pasada hacia delante del entrenamiento sobre la misma longitud.

La caché se rompe además en el límite que impone la cuarta decisión. Pasados los 1024 tokens, generate recorta la secuencia por la izquierda. A partir del segundo bloque, todos los vectores key y value guardados se calcularon cuando el token descartado aún estaba en la secuencia, y como las posiciones son aprendidas, cada token pasa también a otra fila de wpe. La caché entera queda obsoleta. El bucle de nanoGPT se libra porque de todas formas lo recalcula todo, pero un bucle con caché tiene que resolverlo.

Mueve el deslizador de tokens nuevos y verás cómo se separan los dos recuentos. Después elige XL en bf16 y sube el batch hasta que la barra de la caché supere a la de los pesos.

Mismos pesos, bugs distintos

Dos programas que usan los mismos pesos deberían coincidir, y los míos han coincidido, o casi. En 200 pasos de decodificación voraz (greedy), en los que se toma siempre el token con la puntuación más alta, los dos bucles han escrito el mismo texto, aunque sus logits han llegado a diferir en 0,011. Casi toda esa diferencia era un desplazamiento común a todos los tokens, que la softmax ignora. La distancia entre los dos mejores candidatos ha variado como mucho 0,00005, y como en ningún paso ha bajado de 0,004, los tokens han aguantado. Una ejecución más larga o un modelo más grande pueden comerse ese margen, y cuando pase no saltará ningún error.

La diferencia viene de la coma flotante. Pasado el prompt, cada multiplicación de matrices dentro de los bloques del bucle con caché trabaja con una sola fila, mientras que sin caché ve la secuencia entera. Por eso la biblioteca elige otros kernels (las rutinas que hacen las cuentas), que suman en otro orden. Como la suma en coma flotante no es asociativa, los resultados difieren en los últimos bits.

Cambiemos ahora algo que en un despliegue se toca a menudo. He pasado los mismos pesos a bf16, el formato de 16 bits de la parte 3, y he vuelto a ejecutar el bucle de decodificación voraz. El modelo en fp32 ha continuado el prompt con " were not working." y el de bf16 con " were not being run on the same system". Se han separado en el tercer token y ya no han vuelto a juntarse. Al ejecutar el prompt ocho veces en un mismo batch, los tokens se han mantenido, pero los logits se han desplazado hasta 0,005 respecto a la ejecución con un solo prompt.

El remedio al que suele recurrir la gente es poner la temperatura a 0. No lo arregla, y en nanoGPT el 0 ni siquiera es un valor válido. generate divide los logits entre la temperatura. Con 0, los 50 257 logits se vuelven infinitos, la softmax los convierte en NaN (el valor de coma flotante para un resultado indefinido) y torch.multinomial lanza una excepción. La mayoría de las API que aceptan el 0 lo usan como interruptor para pasar a argmax, que elige la puntuación más alta. Eso vuelve a ser decodificación voraz, y solo es tan determinista como lo sean los pesos, el dtype, la forma del batch y los kernels que tenga debajo.

Para ser justos con la magnitud del problema, todas las ejecuciones que aparecen aquí han sido repetibles; con la misma semilla en la misma máquina, el muestreo ha sacado dos veces los mismos 40 tokens. La deriva aparece cuando cambia algo por debajo, y en la infraestructura que sirve el modelo esas cosas cambian sin consultarte.

El muestreo es una política por encima del modelo

La decodificación voraz es una elección, y GPT-2 small enseña por qué se suele evitar: antes de llegar a 200 tokens, el bucle de arriba ya repetía su propia frase. El muestreo, en cambio, sortea el token según la distribución, y tres ajustes dan forma a esa distribución antes del sorteo.

La temperatura divide los logits antes de la softmax, de modo que por debajo de 1 concentra la distribución y por encima de 1 la aplana. Top-k se queda con los k tokens de puntuación más alta y descarta el resto; top-p, también llamado muestreo de núcleo (nucleus sampling), conserva el conjunto más pequeño cuyas probabilidades suman al menos p. El sample.py de nanoGPT trae por defecto temperatura 0,8 y top-k 200, y no tiene top-p.

El efecto de estos ajustes depende del prompt. Tras «The capital of France is», el candidato favorito de GPT-2 small es " the", con un 8,5 %, y " Paris" queda quinto, con un 3,2 %. A temperatura 1, hacen falta 1504 tokens para cubrir el 90 % de la probabilidad. Tras «import numpy as», en cambio, " np" se lleva él solo el 85 %, y para llegar al 90 % bastan tres. La entropía, medida en bits, le pone número a esa dispersión. Para el prompt de la capital a temperatura 1 vale 8,7 bits, tanta incertidumbre como elegir al azar entre unos 400 tokens igual de probables. Para el de numpy vale 1,5 bits, el equivalente a unos tres.

Top-k no distingue un caso del otro, de ahí que cualquier k que elijas falle en uno de los dos. Con los valores por defecto de nanoGPT, esos 200 tokens reúnen el 89 % de la masa tras el prompt de la capital, y tras el de numpy conservan 199 opciones que se reparten menos del 5 %. Top-p sí se adapta, y con la misma temperatura y p = 0.9 conserva 237 tokens en un caso y uno solo en el otro.

Con los valores iniciales, temperatura 1 y p = 0.9, el prompt de la capital conserva 1504 tokens y el de numpy, tres. Sube la temperatura a 2 y verás cómo crece el recuento de numpy. Los indicadores de top-40 y top-200 muestran lo que conserva un k fijo.

Qué vigila el ingeniero de ML

Son cuatro hábitos, y el primero es un test.

Compara en un test el bucle con caché y la pasada completa, logits incluidos. Ejecuta los dos programas durante unos cientos de pasos y compara los tokens de la decodificación voraz y, con una tolerancia, los logits. Si solo compruebas que los tokens coinciden, el test puede pasar mientras los logits derivan, como en mi caso.

Incluye la caché en el presupuesto de memoria, igual que los pesos. Por cada secuencia guarda 2·L·T·d números, con L bloques, T posiciones y anchura d: un vector key y otro value por cada bloque y posición. En GPT-2 small, con 1024 posiciones en fp32, son 75 MB. En GPT-2 XL con bf16 son 315 MB por secuencia, así que unas diez secuencias de longitud completa pesan tanto como sus 3,1 GB de pesos. Subir el tamaño de batch es una decisión de throughput y también de memoria.

Fija lo que comparas. Cualquier cambio de dtype, de forma del batch, de versión de la biblioteca o de kernel altera los logits. Si tu test de regresión compara salidas del modelo, fija los cuatro, o compara con una tolerancia y registra lo que ha cambiado.

Registra los ajustes de muestreo con cada salida. La temperatura, top-k, top-p y la semilla forman parte del programa. Dos ejecuciones que solo se diferencian en ellos se comportan de forma distinta, y una salida registrada sin ellos no se puede reproducir.

Esto tampoco es trabajo de modelado. Los pesos son un único artefacto, y es el código que los rodea el que decide si una posición se puntúa o se escribe, qué sobrevive entre pasos y cómo una distribución se convierte en un token. Todo eso es código corriente que puedes probar.