Entrenar un LLM pequeño desde cero en un portátil con 8 GB de RAM

Mucha gente da por hecho que hace falta un rack de GPU para entrenar un modelo de lenguaje desde cero. No es así, al menos si el objetivo es aprender. Un portátil gaming económico con un Intel i5-9300H, 8 GB de RAM y una GTX 1650 con 4 GB de VRAM basta para construir un transformer pequeño, alimentarlo con un dataset de juguete y ver cómo mejora. Irá lento y algo justo de recursos. También te enseñará más que cualquier API alojada.

Aquí va la primera fase de un proyecto así: el pipeline de datos. El modelo llegará después, así que por ahora el trabajo consiste en tokenizadores, datasets, límites de memoria y lotes.

Por qué merece la pena

Hay tres buenas razones para meterse en esto:

• El Python se oxida cuando un asistente lo escribe todo por ti, y recuperar el hábito a mano es un buen remedio.
• Los sistemas detrás del bombo actual parecen menos misteriosos cuando has montado uno pieza a pieza.
• Implementar artículos por tu cuenta es el camino más corto para leer investigación con espíritu crítico y, quizá algún día, producir la tuya.

El proyecto se llama Axiom, por la palabra que designa una regla aceptada como verdadera sin demostración. Encaja con un primer paso: escribir un transformer uno mismo y luego entrenarlo, partiendo de lo más básico.

Seguir una clase sin copiarla

La serie de vídeos de Andrej Karpathy sobre cómo construir GPT es la guía natural. Su clase empieza con un tokenizador a nivel de carácter. Si ya entiendes la tokenización, puedes adelantarte y usar byte-pair encoding (BPE) desde el principio. Se trata de aprender, así que tiene sentido elegir un camino algo más difícil que el de la pantalla y no volver a teclear su implementación tal cual.

Un proyecto, no un notebook

Los notebooks son cómodos para experimentar, pero difuminan la estructura. Quien vive en Neovim tiene una razón más para dejarlos. Monta en su lugar un proyecto Python normal, con manejo de errores, un cargador de datos y cada paso como una función con nombre. Cada celda del notebook se convierte en una función, y un pequeño orquestador las llama en orden.

Unas cuantas decisiones lo mantienen ordenado:

1. Un directorio .local guarda todo lo pesado que no debe subirse al repositorio: datasets en bruto, tokenizadores y archivos de tokens ya convertidos. Su objetivo es imitar lo mejor de una sesión de notebook, donde no hay que recargar nada en cada ejecución.
2. Un único archivo de configuración recoge los ajustes del modelo, las rutas y los nombres de los tokenizadores, así que nada queda fijado a mano en el pipeline.
3. Cada paso de descarga comprueba si el archivo ya existe antes de volver a bajarlo.

Datos y tokenizador

Para el texto se usa el dataset roneneldan/TinyStories de Hugging Face. Es pequeño sin ser ridículamente diminuto, y su vocabulario sencillo encaja con un modelo que tiene que aprender con poca capacidad de cálculo. Lo sugirió un asistente, así que la elección no se investigó a fondo, pero es razonable para una primera ejecución. Una función de descarga lo baja una vez y lo guarda en caché.

Con el tokenizador pasa lo mismo. Aquí es lakhera2023/gemma4-nano-tinystories, cargado con otro pequeño descargador condicional. De nuevo, una sugerencia aceptada por confianza y luego programada a mano, que es una forma honesta de aprender: deja que el asistente proponga y escribe tú el código para entender qué hace.

El muro de la memoria

En la tokenización es donde una máquina de 8 GB se resiste. El enfoque obvio es tokenizar cada historia, añadir los ID a un único array en memoria que no para de crecer y escribir el array al terminar. En un portátil con poca RAM, ese enfoque congela todo el sistema.

La razón es sencilla cuando la miras. Un token se guarda como entero, y una lista de enteros ocupa mucho más espacio que el texto plano del que procede. La sobrecarga de las listas de Python lo empeora, y ampliar una y otra vez una estructura grande va fragmentando la memoria. Millones de historias después, al sistema operativo no le queda nada.

La solución es no retener la salida en absoluto. Escribe los tokens de cada historia directamente en un archivo en disco en cuanto se generan:

with open(output_file, "wb") as f:
        for sample in dataset:
            text = sample["text"]
            tokens = tokenizer.encode(text)
            tokens.append(tokenizer.eos_id())
            np.asarray(tokens, dtype=np.uint16).tofile(f)

Ahora el uso de memoria es plano, por mucho que crezca el dataset. Los tokens acaban en un archivo binario que luego se puede mapear de vuelta a memoria, que es también como funcionan los pipelines de entrenamiento más grandes. Es el único cuello de botella real de esta fase, y descubrirlo a base de cuelgues es una lección que no se olvida.

Construir los lotes

El batching es más sencillo de lo que parece, aunque requiere un pequeño calentamiento mental. Un modelo de lenguaje aprende a predecir el siguiente token, así que el objetivo de cualquier entrada es esa misma secuencia desplazada una posición hacia delante. Toma un tramo de tokens como x, el tramo que empieza una posición después como y, y apila muchos pares así en un lote. Para cada posición de x, el modelo ve en y lo que viene a continuación.

Esa relación de desplazamiento en uno es todo el truco detrás de la señal de entrenamiento. Parece casi demasiado fácil, y por eso vale la pena implementarla a mano una vez en lugar de llamar a una función de biblioteca.

Lo que viene después

La parte de datos está lista: un dataset en caché, un tokenizador en caché, un archivo de tokens en disco y un generador de lotes. La siguiente fase es el modelo en sí, un transformer escrito desde cero, siguiendo todavía la clase de Karpathy y desviándose cuando haya motivo. Un portátil tan pequeño limitará el tamaño del modelo y la longitud de contexto, así que el verdadero trabajo será meter algo útil en 4 GB de VRAM.