Construye GPT-2 desde cero con PyTorch

Construir tu propio GPT-2 desde cero en Python con PyTorch deja de ser un misterio cuando entiendes cómo cada bloque de la arquitectura se traduce en código ejecutable. Aquí verás, paso a paso, cómo armar el modelo en Google Colab usando GPU, cargar los pesos preentrenados y generar texto real.

Esta guía es para quienes ya conocen la teoría de transformers y quieren ver la implementación concreta: capas de atención, feed forward, layer norm y el ensamble final del modelo de 124 millones de parámetros aproximadamente.

Cómo preparar el entorno en Google Colab con GPU T4

Antes de escribir código, necesitas el hardware correcto. En Colab, vas a Recursos, luego a Cambiar tipo de entorno de ejecución y seleccionas GPU T4 [0:14]. Esto acelera el entrenamiento y la inferencia frente a una CPU.

Luego instalas tres librerías base: annox, xformers y np, que aportan utilidades para programar transformers [0:38]. Después importas torch y numpy, junto con módulos específicos como LayerNorm, que ya viste en la arquitectura de GPT-2.

Para confirmar que estás usando la tarjeta gráfica, ejecutas torch.device y debe responder CUDA [1:18]. Si te aparece cpu, revisa la configuración del entorno.

¿Qué es CUDA en PyTorch? Es la plataforma de NVIDIA que permite a PyTorch ejecutar operaciones sobre la GPU en lugar de la CPU. Si tu device es CUDA, tu modelo correrá mucho más rápido.

Qué bloques componen la arquitectura de GPT-2

GPT-2 se construye apilando piezas que ya viste en la teoría: una red feed forward, un bloque de atención multi-head, normalizaciones y la repetición del transformer 12 veces.

Cómo se programa la feed forward

La feed forward es una red neuronal clásica con dos capas de convolución, una función de activación GELU y un dropout [2:22]. El dropout existe para que la red no sobreaprenda y no memorice solo los datos de entrenamiento; aquí se inicializa en 0.1, un valor bajo pero efectivo.

Dentro de cualquier módulo de PyTorch siempre defines una función forward, que es la que se ejecuta al llamar al módulo. La función backward, que calcula derivadas para el entrenamiento, la genera PyTorch automáticamente gracias a nn.Module [3:01].

También aparece la convolución 1D, una pieza más de ingeniería que optimiza el uso de la tarjeta gráfica. Notarás que hay una conv1d al inicio y otra al final del bloque de atención: la primera transforma la matriz a una forma específica y la segunda la regresa a su forma original [6:43].

Cómo funciona el bloque de atención multi head

Este es el corazón de GPT-2. La clase de atención recibe varios parámetros clave:

• dmodel: dimensión de entrada, en este caso 768, que corresponde al tamaño del embedding [4:31].
• head: cantidad de cabezas de atención, por default 12 [5:00].
• d_head: dimensión de salida de cada cabeza, en este caso 64 [5:32].
• bias: booleano que decide si aplicar sesgo.
• scale: booleano que decide si dividir por la raíz cuadrada de d sub k para estabilidad numérica [5:42].

La función _attention recibe los tres vectores que ya conoces: Q (query), K (keys) y V (values). El primer paso es multiplicar Q por la transpuesta de K usando torch.matmul [8:25]. La transpuesta solo cambia la forma de la matriz: una matriz uno por dos pasa a dos por uno.

Después aplicas la función softmax, que comprime valores grandes a un rango entre -1 y 1, luego el dropout, y al final multiplicas por V para obtener la salida [10:00]. Ese resultado te dice cómo modificar los embeddings para guardar información contextual.

¿Para qué sirve splitheads en multi head attention? Divide el trabajo de las cabezas en batches para paralelizarlo en la GPU. Si tienes que procesar mucha información y no cabe en una sola tarjeta, repartes la carga entre varios procesos.

Cómo ensamblar el transformer block y el modelo completo

El transformer block es la combinación más sencilla del proyecto. Recibe los parámetros, instancia la atención con dmodel=768 y bias=True, agrega la feed forward y dos capas de normalización [11:30].

En la función forward ves cómo encaja todo el diagrama: el input X entra a la multi-head attention, se le aplica un layer norm, luego una feed forward y otra normalización. La teoría visual de la clase anterior se traduce línea por línea aquí.

Por qué GPT-2 usa 12 transformers apilados

El modelo GPT-2 recibe estos parámetros principales:

• layers: número de transformers apilados, en este caso 12 [12:34].
• contexto: ventana de 1024 tokens, que define cuántos tokens del pasado puede ver el modelo para entender una palabra [13:08].
• dmodel: 768 dimensiones por embedding.
• vocabulary_size: 50.257 palabras distintas en el dataset de entrenamiento [13:31].

La función get_clones que creaste al inicio te permite duplicar el bloque transformer 12 veces sin escribirlo a mano. Después se añaden los embeddings WTE y WPE (que traducen la respuesta de las atenciones), un dropout, una layer norm y una capa lineal final que decide cuál de las 50.257 palabras es la siguiente [14:00].

La función de pérdida es la que entrena el modelo: calcula qué tan lejos está la predicción del valor real y actualiza las matrices de pesos de queries, keys y values. La función init_weights inicializa esas matrices con una distribución uniforme [14:46].

¿Qué es la ventana de contexto en GPT-2? Es el máximo de tokens pasados que el modelo puede leer para predecir el siguiente. GPT-2 fue entrenado para 1024 tokens; GPT-3 amplía esto y suma más capas (96 en lugar de 12).

Cómo cargar pesos preentrenados y generar texto

Una vez programado el modelo, importas el GPT Tokenizer y descargas los pesos preentrenados de GPT-2, que son de código abierto [15:38]. Cargas esos pesos binarios al modelo, actualizas los estados y verificas que la arquitectura compile correctamente: deberías ver el module list con los 12 transformers, los embeddings, el dropout, la layer norm, la capa lineal y la cross entropy.

Para medir el tamaño del modelo, sumas los parámetros con num_elements y los traduces a bytes asumiendo float16 (4 bytes). El resultado es 474 MB, lo bastante ligero para correr incluso en un teléfono móvil [16:48].

La generación de texto usa un for loop dentro de la función generate: toma un token, lo ejecuta, lo agrega al contexto y vuelve a ejecutar. Con el prompt The planet Earth is y 40 tokens, el modelo tarda 7.6 segundos en GPU T4 y produce una continuación coherente como The planet Earth is a beautiful place that... [17:52].

Recuerda que cada token no equivale a una palabra completa: a veces son subwords. Y el modelo está entrenado en inglés, así que los resultados en español serán pobres.

Ahora te toca a ti. Cambia el número de capas a 96 para simular GPT-3, ajusta el dropout, prueba activar scale=True o aumenta los tokens generados. ¿Qué resultados obtuviste? Cuéntame en los comentarios cómo modificarías la máscara de atención que dejé como Easter egg.