Proyecto: ajustar un modelo pequeño en Colab

De la cadena que el bloque ha montado sobre el mini-GPT, el primer eslabón dejó una predicción por comprobar. El ajuste supervisadoBloque 2 · Lección 3Ajuste supervisado: máxima verosimilitud sobre un formatoSeguir entrenando un modelo preentrenado sobre pares de instrucción y respuesta es la entropía cruzada de siempre, cobrada sólo en la respuesta: máxima verosimilitud sobre el conjunto de pares. La lección deriva su óptimo, que copia el conjunto, y su gradiente, que es el del preentrenamiento con las filas del prompt a cero. Después ajusta el mini-GPT con diez preguntas: aprende el formato, lee la pregunta, inventa el contenido y olvida su novela. (supervised fine-tuning, SFT) le enseñó a un modelo preentrenado con $125\,789$ tokens de Galdós la forma de diez respuestas y ningún contenido, y la lección anticipó lo que pasaría en el otro extremo: con miles de pares variados y un modelo que lo ha leído casi todo, respuestas útiles. De ese otro extremo salen los modelos abiertos, y la referencia de la que parte todo lo que el bloque ha puesto después.

¿Qué cuesta comprobarlo? Esa cuenta es lo primero que hace esta lección, porque decide dónde corre el proyecto. El modelo es Qwen2.5-0.5B, uno de los modelos abiertos más pequeños que se publican, y ajustarlo entero pide casi el doble de la memoria que el Python del navegador puede llegar a usar: su implementación es un cuaderno de Colab, con una GPU (graphics processing unit) prestada. Lo segundo es mirar dentro del modelo antes de ajustarlo, y lo que aparece es un problema que el bloque ya conoce.

El cuaderno repite, línea por línea, el ajuste supervisado del mini-GPT: una plantilla con otros nombres, la misma máscara, la misma pérdida y el bucle de la lección sobre entrenar el mini-GPTBloque 1 · Lección 3Entrenar un mini-GPT en NumPyEl checkpoint del bloque es el resultado de 2 000 pasos de un bucle que esta lección escribe entero. Cada paso sigue la pérdida media de unas pocas ventanas del corpus sacadas al azar, y la lección deriva por qué eso apunta, de media, adonde apuntaría el corpus entero, y cuánto ruido cuesta ahorrarse el resto. El paso es AdamW, con una tasa de aprendizaje que se calienta y luego decae. Con todo ello el mini-GPT sigue entrenando en el navegador, sobre las últimas páginas de Marianela, que no había visto.. Cambian los números:

La fila que decide el resultado es la tercera. La pérdida del ajuste sólo habla de los pares del conjunto, y lo que el modelo haga con cualquier otra instrucción lo deciden los pesos que comparten todas y el sitio del que partieron: es lo que la pérdida no diceBloque 2 · Lección 3Ajuste supervisado: máxima verosimilitud sobre un formatoSeguir entrenando un modelo preentrenado sobre pares de instrucción y respuesta es la entropía cruzada de siempre, cobrada sólo en la respuesta: máxima verosimilitud sobre el conjunto de pares. La lección deriva su óptimo, que copia el conjunto, y su gradiente, que es el del preentrenamiento con las filas del prompt a cero. Después ajusta el mini-GPT con diez preguntas: aprende el formato, lee la pregunta, inventa el contenido y olvida su novela.. Para el mini-GPT, ese sitio era una novela, y diez pares de colores no tenían ningún contenido que sacar de ella. Qwen leyó en su preentrenamiento unos ciento cuarenta millones de veces más texto, en más de 29 idiomas, el español entre ellos. El contenido ya está en sus pesos, y lo que los pares tienen que enseñarle es a sacarlo en forma de respuesta.

Lo que cuesta ajustar 494 millones de parámetros

Contemos primero lo que hay que mover. Qwen2.5-0.5B tiene $L = 24$ bloques de $d_{\text{model}} = 896$ y un vocabulario de $\lvert V \rvert = 151\,936$ entradas, y el fichero de sus pesos trae $N = 494\,032\,768$ parámetros. Más de la cuarta parte son la tabla de embeddings, atada a la salida como la del mini-GPT:

El total sale del fichero y no de la cuenta por bloque de la lección sobre la cachéBloque 1 · Lección 5La caché de claves y valores: cuánto cuesta cada tokenPara escribir cada token, el mini-GPT volvía a pasar el texto entero por la red. La lección demuestra que las claves y los valores de las posiciones ya leídas no cambian, los guarda en una caché y cuenta lo que cuesta cada token con ella y sin ella: sin ella, exactamente t veces más; con ella, una memoria que crece con el texto y que en un modelo grande llega a ocupar más que sus pesos. Todo medido sobre el mini-GPT, y contra el reloj., que aquí no da exacta: el perceptrón de Qwen tiene tres matrices y no dos, y sus cabezas comparten las claves y los valores por grupos.

Entrenar pide más que los pesos. Cada parámetro lleva al lado su gradiente, y AdamW guarda por cada uno las dos medias móviles del Adam del mini-GPT, $\mathbf{m}_s$ y $\mathbf{v}_s$. Son cuatro números por parámetro, de cuatro bytes cada uno en coma flotante de 32 bits:

y todavía no se ha guardado ni una activación. Las de la ida, que la vuelta necesita, crecen con la longitud de la conversación, y sólo los logits de una posición son ya $151\,936$ números: con 512 posiciones, $0.31$ GB.

Una pestaña del navegador no llega. Pyodide, el Python que ejecuta las celdas del curso, está compilado a WebAssembly de 32 bits: una dirección de memoria tiene 32 bits, y con ellas se numeran como mucho $2^{32}$ bytes, unos $4.3$ GB. Los pesos y sus gradientes ya ocupan $3.95$ de ellos. A Pyodide le falta también PyTorch, pero eso tendría arreglo, porque el mini-GPT se entrenó en NumPy sin él. La memoria, no.

Ni el tiempo. La lección sobre entrenar el mini-GPT midió en el navegador algo menos de un milisegundo por token de ida y vuelta para sus $136\,448$ parámetros, y lo que cuesta un token es proporcional a $N$: es el $6N$ de la lección sobre las leyes de escalaBloque 1 · Lección 7Leyes de escala: Kaplan y ChinchillaLa pérdida de un modelo de lenguaje baja como una potencia de sus parámetros, de los tokens que lee y del cálculo que gasta: una recta en ejes logarítmicos. La lección cuenta lo que cuesta entrenar, C ≈ 6ND, deriva el reparto óptimo de un presupuesto entre modelo y datos, lee en las cifras de Chinchilla sus 20 tokens por parámetro y hace las cuentas del mini-GPT: su tamaño es el que Chinchilla le daría a su cálculo, y lo que le falta es texto.. Qwen tiene $3\,621$ veces más parámetros, unos $3.6$ segundos por token. Los pares del cuaderno suman $788\,511$ tokens, y una sola época tardaría

más de un mes con la pestaña abierta. Las celdas del curso se cortan a los diez segundos.

En una GPU, las cuentas cambian de escala. El cálculo de esa época, en operaciones de coma flotante (floating-point operations, FLOPs), es el $C \approx 6ND$ de aquella lección con $D = 788\,511$:

La T4 que Colab presta gratis hace hasta $8.1 \times 10^{12}$ de esas operaciones por segundo en coma flotante de 32 bits. A pleno rendimiento son cinco minutos; en la práctica, más. Tiene 16 GB de memoria, y el ajuste cabe entero: los $7.9$ GB, y al lado las activaciones de una conversación cada vez.

Si no cupiera, la salida habitual sería no mover los $N$ parámetros. LoRA (low-rank adaptation) los congela y entrena, al lado de cada matriz, dos matrices estrechas cuyo producto se le suma, y los gradientes y las medias móviles se quedan en una fracción mínima de los $7.9$ GB. Aquí no hace falta, y queda fuera del curso.

La misma pérdida, con la plantilla de Qwen

Qwen usa ChatML (Chat Markup Language), la plantilla que la lección sobre la plantilla de chatBloque 2 · Lección 2La plantilla de chat es un formato entrenadoUn modelo lee una sola secuencia de tokens, y una conversación es una lista de mensajes con autor. La plantilla de chat escribe la una como la otra: cada mensaje entre el token especial de su rol y el de fin, que ningún texto puede producir, y eso la hace inyectiva, cosa que unas marcas de texto no consiguen. Dentro de ella la respuesta tiene un sitio exacto, su probabilidad es un producto y la máscara dice qué tokens se cobran al entrenar. Y un token especial que nadie ha entrenado no significa nada: la lección acota cuánto puede valer y lo mide en el mini-GPT, que no sabe terminar su respuesta. nombró como la más extendida entre los modelos abiertos. Tiene dos tokens especiales y una diferencia con la del mini-GPT: el rol no tiene token propio. Cada mensaje abre con <|im_start|>, sigue con el nombre del rol escrito en tokens corrientes, user o assistant, y un salto de línea, y cierra con <|im_end|>, el token de fin de mensaje, y otro salto.

La máscara es la de aquella lección: 1 en el contenido de la respuesta y en su <|im_end|>, 0 en todo lo demás, el rol del asistente y su salto incluidos, porque los escribe la plantilla antes de que el modelo tome la palabra. En la respuesta de la nieve se cobran ocho tokens, los siete de La nieve es blanca. y el <|im_end|>. La pérdida es la $\mathcal{L}_{\text{SFT}}$ de la lección sobre el ajuste supervisado sin cambiar un símbolo: en cada paso, la suma de $-\log p_\theta$ sobre los tokens cobrados de un batch, dividida entre cuántos son.

Dos filas que el preentrenamiento no distinguió

Ahí está el problema conocido. En el modelo preentrenado, la fila de <|im_end|> en la tabla de embeddings es, hasta el redondeo, la de <|im_start|>, y la de las 290 entradas que van detrás hasta el final de la tabla: otros especiales que Qwen reserva para otros usos y filas de relleno que no corresponden a ningún token. El cuaderno lo comprueba: las 292 filas miden entre $0.3009$ y $0.3010$, y el coseno entre dos cualesquiera es 1 hasta la cuarta cifra. El preentrenamiento no distinguió nunca entre ellas.

Con la tabla atada, el logit de una entrada es $z_v = \mathbf{e}_v^{\top}\mathbf{h}$, y filas iguales dan logits iguales en cualquier posición. Las 292 entradas empatan, y el softmax les da a todas la misma probabilidad. Entre todas no pasan de 1, así que cada una, y <|im_end|> con ellas, se queda en

en cada posición de cualquier respuesta. Es la cota de la lección sobre la plantilla, $1/516$ para un especial del mini-GPT, con otra razón detrás: allí, las filas nuevas empezaban en la media de la tabla; aquí, el preentrenamiento las ha dejado iguales.

La consecuencia se ve en cuanto el modelo escribe. Si cada posición cierra con probabilidad $1/292$ como mucho, una respuesta muestreada tiene de media 292 tokens o más, y la generación voraz, que sólo elegiría <|im_end|> si ninguna otra entrada pasara de $1/292$, no la cierra. El modelo preentrenado no sabe terminar. De todo lo que el ajuste le enseña, lo más nuevo es el token que cierra cada respuesta, y ningún texto se lo enseñó antes.

El cuaderno, y qué mirar mientras corre

El proyecto no corre sobre tu mini-GPT ni sobre ningún modelo del navegador. Corre sobre Qwen2.5-0.5B en su versión preentrenada, la que Qwen publica como base, frente a la instruct, ya ajustada. Su licencia es Apache 2.0, y el cuaderno lo carga fijado a un commit, igual que fija las versiones de las bibliotecas.

Antes de ejecutar nada, elige en el menú de Colab un entorno con GPU; la segunda celda dice cuál te ha tocado. Después, el cuaderno recorre esta lección en orden. Imprime las cuentas de la primera sección, escribe la plantilla a mano y la compara con la que trae la biblioteca (que, si no le das un mensaje de sistema, añade uno en inglés que nadie ha escrito), y mide las 292 filas. Le hace al modelo preentrenado cinco instrucciones que he escrito yo, y ninguna de las cinco respuestas llega a su <|im_end|>: cuatro repiten la instrucción una y otra vez, y la quinta da vueltas a una frase.

Los pares salen de OpenAssistant Conversations, un conjunto que escribieron y puntuaron más de $13\,500$ voluntarios en 2023, en 35 idiomas, con licencia Apache 2.0. De cada conversación en español, el cuaderno toma el primer mensaje y la respuesta que los voluntarios pusieron primera al compararla con las demás: un par de ajuste supervisado escogido con una comparación, como las de la lección sobre de dónde salen las preferenciasBloque 2 · Lección 7De dónde salen las preferencias: humanos, constituciones y modelosRLHF y DPO no juzgan sus comparaciones: las reproducen. Quién las escribe decide lo que la política prefiere, y hay tres fuentes: personas con una guía, un modelo que critica y revisa sus respuestas según una constitución, y un modelo que juzga. En papel, una mayoría del 60 % pasa a escribirse el 98 % de las veces con β = 0.1. «Alineado» quiere decir alineado con alguien, y no quiere decir correcto.. Descarta las conversaciones de más de 512 tokens en vez de cortarlas, porque una respuesta cortada no llega a su <|im_end|>. Quedan $3\,249$ pares con $788\,511$ tokens, de los que se cobra el 86 %, y los 169 de la parte de validación quedan reservados para medir.

El ajuste es una época, 203 pasos de 16 conversaciones, con el bucle de la lección sobre entrenar el mini-GPT: AdamW, recorte del gradiente a norma 1 y su calendario, calentamiento y coseno hasta la décima parte. La tasa máxima es $\eta_{\max} = 2 \times 10^{-5}$, 250 veces menos que la del ajuste del mini-GPT, que tenía que verse en doce pasos y se llevó la novela por delante; ésta tiene que dejar en su sitio lo que Qwen sabe. Las conversaciones pasan de una en una y cada una suma su gradiente al del batch, así que ninguna se rellena. Al final, la celda imprime cuánto ha tardado y la memoria que llegó a ocupar la GPU: compárala con los $7.9$ GB.

El cuaderno mide tres números antes del ajuste y después. La pérdida en las respuestas reservadas, en nats por token cobrado: antes, $1.75$, cuando repartir por igual costaría $\ln 151\,936 \approx 11.9$. El texto de las respuestas ya lo predice bien, aunque no sepa escribirlo en su turno. La probabilidad media que da a <|im_end|> justo donde una respuesta reservada acaba: antes, $5.1 \times 10^{-8}$, más de sesenta mil veces por debajo de la cota. Y la pérdida en el final de Marianela, de donde salió el trozo con que la lección sobre el ajuste supervisado midió el olvido: antes, $3.24$ nats por token de Qwen, un número que sólo se compara con el suyo de después. Mira cuánto se mueve cada uno, y vuelve a darle las cinco instrucciones. La última es la pregunta de la Gran Muralla de la lección anterior: fíjate en qué contesta, y pregúntate de dónde lo saca.

Comprueba tu intuición

Tres preguntas: por qué esto no cabe en el navegador, qué tokens cobra el ajuste y qué asegura la cota.

El ajuste supervisado hace obedecer a un modelo en la forma de sus pares. Lo que el modelo dice, en cambio, sale de dos sitios y de ninguno más: de lo que sus pesos guardaron del preentrenamiento y de cómo escribían quienes redactaron las respuestas. Nada en el ajuste comprueba que sea verdad. Si los pesos no saben si la Muralla se ve desde la Luna, la respuesta tendrá igual forma de respuesta.

Queda una palanca que no toca los pesos: lo que va en el prompt. Un prompt puede darle al modelo lo que sus pesos no saben, pedirle una respuesta que se pueda comprobar y decirle con exactitud qué se espera de ella. La lección siguiente abre el bloque 3 por ahí, con el prompt como especificación.