Entrena tu propio LLM desde cero: Una guía completa de principio a fin
Un tutorial práctico para construir, preentrenar y alinear un Transformer LLM desde cero usando PyTorch, cubriendo SFT, DPO, PPO y GRPO.
¿Alguna vez has querido entrenar tu propio modelo de lenguaje grande (LLM) desde cero, no solo ajustar uno existente? El proyecto de código abierto train-llm-from-scratch de FareedKhan-dev proporciona un pipeline completo y transparente que te lleva desde texto sin procesar hasta un modelo alineado y capaz de razonar. Esto no es un envoltorio de caja negra alrededor de bibliotecas existentes; es un viaje educativo práctico donde cada algoritmo—desde la atención multi-cabeza hasta GRPO—está implementado en PyTorch puro.
Esta guía te llevará a través de todo el proceso, explicando el por qué y el cómo detrás de cada paso, para que puedas entender los fundamentos y ejecutar el código tú mismo.
¿Por qué construir un LLM desde cero?
La mayoría de los desarrolladores interactúan con LLMs a través de APIs de alto nivel como transformers de Hugging Face o trl. Aunque son poderosos, abstraen los mecanismos centrales. Construir desde cero te da:
- Comprensión profunda: Aprendes exactamente cómo funcionan la tokenización, la atención y los bucles de entrenamiento.
- Control total: Puedes modificar cualquier componente—funciones de pérdida, arquitecturas o estrategias de entrenamiento—sin restricciones de bibliotecas.
- Valor educativo: Es la mejor manera de comprender conceptos como el enmascaramiento causal, el modelado de recompensas y los métodos de gradiente de política.
Este proyecto está diseñado para estudiantes, desarrolladores e investigadores que quieren ver todo el pipeline en un solo lugar, desde la descarga de datos hasta la generación de texto.
El Pipeline de un Vistazo
El viaje sigue un camino claro y secuencial:
- Preparación de Datos: Texto sin procesar → IDs de Token → Arrays almacenados.
- Construcción del Modelo: Construir un Transformer a partir de componentes pequeños y reutilizables (MLP, Atención, Bloques).
- Preentrenamiento: Predicción del siguiente token en un corpus grande (The Pile).
- Generación de Texto: Muestreo del modelo entrenado.
- Post-Entrenamiento (Alineación): Convertir un modelo base en un asistente útil mediante SFT, Modelado de Recompensas, DPO, PPO y GRPO.
- Evaluación: Medir el rendimiento en un benchmark como GSM8K.
- Inferencia y Chat: Interactuar con el modelo final.
Paso 1: Preparando los Datos
Un modelo solo entiende enteros. El primer trabajo es convertir texto en IDs de token y almacenarlos de manera eficiente.
Tokenización
El proyecto utiliza el tokenizador r50k_base de OpenAI (de tiktoken), el mismo usado por GPT-3. El texto se convierte en una lista de enteros, y se añade un token especial <|endoftext|> (ID 50256) a cada documento para que el modelo aprenda los límites.
# Descargar y tokenizar datos de The Pile
python scripts/data_download.py
python scripts/data_preprocess.py
Para la ruta más nueva y rápida:
python scripts/prepare_pretrain_data.py --split val --out data/pile_dev.h5
python scripts/prepare_pretrain_data.py --split train --num_shards 1 --out data/pile_train.h5
Formato de Chat y Enmascaramiento de Pérdida
Para el post-entrenamiento, el modelo necesita entender los roles (usuario vs. asistente). El proyecto utiliza marcadores de texto plano:
<|user|>
¿Cuánto es 13 + 29?<|endoftext|><|assistant|>
<think>13 + 29 = 42</think><answer>42</answer><|endoftext|>
Una máscara de pérdida asegura que el modelo solo aprenda de los tokens del asistente, no del prompt del usuario. Esto es crítico para SFT y RL.
def encode_chat(messages, add_generation_prompt=False):
ids, mask = [], []
for m in messages:
role = m["role"]
header_ids = _encode_ordinary(_header_for(role))
ids.extend(header_ids)
mask.extend([0] * len(header_ids))
content_ids = _encode_ordinary(m["content"])
is_completion = role == "assistant"
ids.extend(content_ids)
mask.extend([1 if is_completion else 0] * len(content_ids))
ids.append(EOT_ID)
mask.append(1 if is_completion else 0)
return ids, mask
Paso 2: Construyendo el Transformer
El Transformer está construido a partir de cuatro módulos pequeños y reutilizables de PyTorch. Este diseño modular hace que el código sea fácil de entender y modificar.
Perceptrón Multicapa (MLP)
La parte "pensante" de cada bloque. Expande el vector del token, aplica una ReLU y lo proyecta de vuelta.
class MLP(nn.Module):
def __init__(self, n_embed):
super().__init__()
self.hidden = nn.Linear(n_embed, 4 * n_embed)
self.relu = nn.ReLU()
self.proj = nn.Linear(4 * n_embed, n_embed)
def forward(self, x):
x = self.relu(self.hidden(x))
x = self.proj(x)
return x
Atención de una Sola Cabeza
Permite que un token mire a otros tokens. Calcula vectores de consulta, clave y valor, aplica una máscara causal (para que un token no pueda ver tokens futuros) y toma una suma ponderada de valores.
class Head(nn.Module):
def __init__(self, head_size, n_embed, context_length):
super().__init__()
self.key = nn.Linear(n_embed, head_size, bias=False)
self.query = nn.Linear(n_embed, head_size, bias=False)
self.value = nn.Linear(n_embed, head_size, bias=False)
self.register_buffer('tril', torch.tril(torch.ones(context_length, context_length)))
def forward(self, x):
B, T, C = x.shape
k = self.key(x)
q = self.query(x)
scale_factor = 1 / math.sqrt(C)
attn_weights = q @ k.transpose(-2, -1) * scale_factor
attn_weights = attn_weights.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
attn_weights = F.softmax(attn_weights, dim=-1)
v = self.value(x)
out = attn_weights @ v
return out
Atención Multi-Cabeza
Ejecuta múltiples cabezas de atención en paralelo, cada una aprendiendo diferentes relaciones, luego concatena y proyecta los resultados.
El Bloque Transformer
Combina atención y MLP con conexiones residuales de pre-normalización. Esta es la unidad repetitiva central.
class Block(nn.Module):
def __init__(self, n_head, n_embed, context_length):
super().__init__()
self.ln1 = nn.LayerNorm(n_embed)
self.attn = MultiHeadAttention(n_head, n_embed, context_length)
self.ln2 = nn.LayerNorm(n_embed)
self.mlp = MLP(n_embed)
def forward(self, x):
x = x + self.attn(self.ln1(x))
x = x + self.mlp(self.ln2(x))
return x
El Transformer Completo
Envuelve todo: embeddings de tokens, embeddings de posición, una pila de bloques, capa de normalización final y una proyección lineal a logits del tamaño del vocabulario.
class Transformer(nn.Module):
def __init__(self, n_head, n_embed, context_length, vocab_size, N_BLOCKS):
super().__init__()
self.token_embed = nn.Embedding(vocab_size, n_embed)
self.position_embed = nn.Embedding(context_length, n_embed)
self.attn_blocks = nn.ModuleList([Block(n_head, n_embed, context_length) for _ in range(N_BLOCKS)])
self.layer_norm = nn.LayerNorm(n_embed)
self.lm_head = nn.Linear(n_embed, vocab_size)
self.register_buffer('pos_idxs', torch.arange(context_length))
def forward(self, idx, targets=None):
x = self.forward_hidden(idx)
logits = self.lm_head(x)
loss = None
if targets is not None:
B, T, C = logits.shape
flat_logits = logits.reshape(B * T, C)
targets = targets.reshape(B * T).long()
loss = F.cross_entropy(flat_logits, targets)
return logits, loss
Paso 3: Preentrenando el Modelo Base
El preentrenamiento es la predicción del siguiente token en un corpus grande. El modelo lee ventanas aleatorias de tokens, predice el siguiente token en cada posición y actualiza sus pesos mediante la pérdida de entropía cruzada.
# Entrenar un modelo de 13M de parámetros
python scripts/train_transformer.py
# Para modelos más grandes con optimizaciones de memoria
python scripts/train_transformer.py --amp --grad-checkpointing --grad-accum 8
La pérdida comienza cerca de ln(vocab_size) (~10.8) y disminuye a medida que el modelo aprende estadísticas del lenguaje. Un modelo de 77M de parámetros entrenado en 2 GPUs L40 alcanzó una pérdida de ~3.73 después de 2000 pasos.
Paso 4: Generando Texto
Una vez entrenado, puedes generar texto muestreando de la distribución de salida del modelo.
def generate(self, idx, max_new_tokens):
for _ in range(max_new_tokens):
idx_cond = idx[:, -self.context_length:]
logits, _ = self(idx_cond)
logits = logits[:, -1, :]
probs = F.softmax(logits, dim=-1)
idx_next = torch.multinomial(probs, num_samples=1)
idx = torch.cat((idx, idx_next), dim=1)
return idx
python scripts/generate_text.py --model_path models/transformer_B.pt --input_text "El" --max_new_tokens 100
Paso 5: Post-Entrenamiento (Alineación)
Un modelo base puede continuar texto, pero no puede seguir instrucciones. El post-entrenamiento lo alinea usando ajuste fino supervisado (SFT) y aprendizaje por refuerzo (RL).
SFT (Ajuste Fino Supervisado)
Enseña al modelo a responder en un formato de chat. La pérdida solo se calcula en los tokens del asistente.
python scripts/prepare_sft_data.py --context_length 1024
torchrun --standalone --nproc_per_node=2 scripts/train_sft.py
El Modelo de Recompensa
Entrena una pequeña cabeza lineal sobre el modelo SFT para puntuar respuestas. Utiliza la pérdida de Bradley-Terry para preferir respuestas elegidas sobre las rechazadas.
def bradley_terry_loss(chosen_rewards, rejected_rewards):
return -F.logsigmoid(chosen_rewards - rejected_rewards).mean()
DPO, ORPO y KTO
La Optimización Directa de Preferencias (DPO) omite el modelo de recompensa y el bucle de RL, trabajando directamente sobre pares de preferencias. ORPO y KTO son variantes.
torchrun --standalone --nproc_per_node=2 scripts/train_dpo.py --loss_type dpo
PPO (Optimización de Política Proximal)
El bucle clásico de RLHF. El modelo genera respuestas, las puntúa con un modelo de recompensa y se actualiza mediante una pérdida de política recortada.
def ppo_policy_loss(new_logp, old_logp, advantages, mask, clip=0.2):
ratio = torch.exp(new_logp - old_logp)
surr1 = ratio * advantages
surr2 = torch.clamp(ratio, 1.0 - clip, 1.0 + clip) * advantages
loss = -masked_mean(torch.min(surr1, surr2), mask)
return loss
GRPO (Optimización de Política Relativa por Grupos)
El método al estilo DeepSeek-R1. Muestrea un grupo de respuestas para cada prompt, las puntúa y usa la media y desviación estándar del grupo como línea base.
def group_advantages(rewards, group_size, eps=1e-4):
r = rewards.view(-1, group_size)
mean = r.mean(dim=1, keepdim=True)
std = r.std(dim=1, keepdim=True)
adv = (r - mean) / (std + eps)
return adv.reshape(-1)
Paso 6: Evaluación
La métrica clave es la precisión greedy en GSM8K: el modelo debe generar la respuesta correcta dentro de las etiquetas <answer>.
for s in base_pretrained sft dpo ppo grpo; do
python scripts/eval_post_training.py --ckpt models/$s.pt --label $s --limit 200 --append logs/table.jsonl
done
Paso 7: Hablando con el Modelo
Usa el script de chat para interactuar con cualquier checkpoint.
python scripts/chat.py --ckpt models/sft.pt --prompt "¿Cuánto es 13 + 29?"
Cómo Empezar
Clona el repositorio:
git clone https://github.com/FareedKhan-dev/train-llm-from-scratch.git cd train-llm-from-scratch pip install -e .Empieza pequeño: Entrena primero el modelo de 13M de parámetros. Es rápido y te dará una idea del pipeline.
Escala: Aumenta
n_embedyn_blockshasta que llegues al límite de memoria de tu GPU.Recorre la cadena de post-entrenamiento: Ejecuta SFT, luego DPO, luego PPO o GRPO, y observa cómo mejora la precisión en GSM8K.
Este proyecto es una joya rara: es tanto una herramienta práctica como un tutorial completo. Ya seas un estudiante que quiere entender los Transformers en profundidad o un desarrollador que busca experimentar con técnicas de alineación, train-llm-from-scratch es el punto de partida perfecto.