CONSTRUYENDO UN LLM DESDE CERO CON PYTORCH

Introducción a los Modelos de Lenguaje Grandes

Los modelos de lenguaje grandes representan un avance significativo en el campo de la inteligencia artificial, permitiendo a las máquinas generar texto coherente y contextualizado similar al humano. En este tutorial, nos sumergiremos en el proceso completo de construcción de un modelo de este tipo utilizando PyTorch, una biblioteca flexible y potente para el desarrollo de redes neuronales. El enfoque será práctico y teórico, cubriendo desde los bloques fundamentales hasta las técnicas de alineación avanzadas que aseguran comportamientos útiles y seguros.

Para contextualizar, considera que en octubre de 2025, la adopción de estos modelos ha evolucionado rápidamente, con integraciones en aplicaciones cotidianas como asistentes virtuales y herramientas de codificación. Actualizando la información disponible, ahora se enfatiza la importancia de la eficiencia computacional debido al aumento en los costos energéticos y la demanda de modelos más sostenibles. Este tutorial te equipará con el conocimiento necesario para experimentar y contribuir a este ecosistema en crecimiento.

Comenzaremos explorando la arquitectura base que sustenta estos modelos. La comprensión de estos componentes es esencial para cualquier desarrollador interesado en la programación de IA. A lo largo del texto, integraremos ejemplos de código que ilustran cada concepto, facilitando la replicación en entornos locales.

Arquitectura Básica del Transformador

La arquitectura del transformador es el pilar sobre el cual se construyen la mayoría de los modelos de lenguaje modernos. Introducida en 2017, esta estructura revoluciona el procesamiento de secuencias al eliminar las dependencias recurrentes en favor de mecanismos de atención paralelizables. En esencia, un transformador consta de codificadores y decodificadores apilados, pero para modelos generativos como los LLMs, nos centramos en la variante decodificadora-only.

El núcleo de esta arquitectura reside en el mecanismo de atención auto-regresiva, que permite al modelo ponderar la importancia de diferentes tokens en la secuencia de entrada. Matemáticamente, la atención se calcula como una función suave de similitud entre consultas, claves y valores derivados de las entradas embebidas. La fórmula clave es: Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V, donde d_k es la dimensión de las claves.

Para implementar esto en PyTorch, inicia con la definición de una clase básica para la atención multi-cabeza. Este componente divide la atención en múltiples “cabezas” para capturar dependencias diversas.

import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        assert d_model % num_heads == 0
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

    def scaled_dot_product_attention(self, Q, K, V, mask=None):
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        attn_weights = torch.softmax(scores, dim=-1)
        output = torch.matmul(attn_weights, V)
        return output

    def forward(self, Q, K, V, mask=None):
        batch_size = Q.size(0)

        Q = self.W_q(Q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(K).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(V).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)

        attn = self.scaled_dot_product_attention(Q, K, V, mask)
        attn = attn.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
        return self.W_o(attn)

Este código define una atención multi-cabeza básica, incorporando una máscara para la generación auto-regresiva, que previene el acceso a tokens futuros. En un modelo completo, esta capa se integra en bloques de transformador que alternan atención y redes feed-forward.

Las capas feed-forward consisten en dos transformaciones lineales con activación ReLU intermedio, aplicadas posición por posición. Su propósito es introducir no linealidades que enriquezcan las representaciones aprendidas. Una implementación simple sería:

class FeedForward(nn.Module):
    def __init__(self, d_model, d_ff):
        super(FeedForward, self).__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.linear2(self.relu(self.linear1(x)))

Combinando estos elementos, un bloque de transformador completo incluye normalización en capas residuales para estabilizar el entrenamiento. La conexión residual se expresa como x + sublayer(LayerNorm(x)), donde sublayer es la atención o feed-forward. Esto mitiga problemas de gradientes vanishing en redes profundas.

En el contexto actual de 2025, optimizaciones como la cuantización de pesos se han vuelto estándar para desplegar estos modelos en dispositivos edge, reduciendo el footprint de memoria sin sacrificar precisión. Experimenta con este código ajustando el número de cabezas para observar impactos en la convergencia.

Entrenando un Modelo Pequeño

Una vez establecida la arquitectura, el siguiente paso es entrenar un modelo diminuto para validar el entendimiento. Elegimos un tamaño reducido, como 128 dimensiones de modelo y 4 cabezas de atención, para que el entrenamiento sea accesible en hardware modesto.

El proceso inicia con la preparación de datos. Utiliza un corpus simple como texto shakesperiano para simplicidad, tokenizado en subpalabras con un vocabulario limitado. PyTorch’s DataLoader facilita el batching y shuffling.

from torch.utils.data import Dataset, DataLoader
import tiktoken  # Para tokenización, asumiendo disponibilidad

class TextDataset(Dataset):
    def __init__(self, text, tokenizer, block_size):
        self.tokenizer = tokenizer
        self.tokens = tokenizer.encode(text)
        self.block_size = block_size

    def __len__(self):
        return len(self.tokens) - self.block_size

    def __getitem__(self, idx):
        chunk = self.tokens[idx:idx + self.block_size + 1]
        return torch.tensor(chunk[:-1]), torch.tensor(chunk[1:])

# Ejemplo de uso
text = "Texto de entrenamiento aquí..."  # Reemplaza con corpus real
tokenizer = tiktoken.get_encoding("gpt2")
dataset = TextDataset(text, tokenizer, block_size=256)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

El bucle de entrenamiento optimiza la pérdida de predicción del siguiente token mediante cross-entropy. Usa AdamW como optimizador, con un learning rate scheduler para ajuste dinámico.

model = TransformerModel(vocab_size=tokenizer.n_vocab, d_model=128, num_heads=4, num_layers=6)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
criterion = nn.CrossEntropyLoss()

for epoch in range(num_epochs):
    for batch_x, batch_y in dataloader:
        optimizer.zero_grad()
        logits = model(batch_x)
        loss = criterion(logits.view(-1, logits.size(-1)), batch_y.view(-1))
        loss.backward()
        optimizer.step()
    print(f"Epoch {epoch}, Loss: {loss.item()}")

Este snippet ilustra un ciclo básico. En práctica, incorpora gradientes clipping para estabilidad: torch.nn.utils.clipgrad_norm(model.parameters(), max_norm=1.0). Al entrenar, monitorea métricas como perplexidad, que mide la incertidumbre del modelo en el corpus.

En actualizaciones recientes, herramientas como DeepSpeed han popularizado el entrenamiento distribuido para modelos pequeños, permitiendo escalado temprano. Este enfoque no solo valida la arquitectura sino que revela desafíos como overfitting, mitigado mediante dropout en el 10-20% de las conexiones.

Explorando más, considera la inicialización de pesos con Xavier para uniformizar la distribución de activaciones, asegurando flujos de gradiente saludables desde el inicio.

Mejoras Modernas en Eficiencia

Las implementaciones básicas de transformadores sufren de ineficiencias en inferencia y entrenamiento a escala. En 2025, técnicas de normalización avanzada como RMSNorm han reemplazado a LayerNorm tradicional por su simplicidad computacional y estabilidad.

RMSNorm normaliza dividiendo por la raíz cuadrada de la media cuadrática de los elementos, sin el término de bias: x_norm = x / sqrt(mean(x^2) + eps). Su implementación es:

class RMSNorm(nn.Module):
    def __init__(self, d_model, eps=1e-6):
        super(RMSNorm, self).__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(d_model))

    def forward(self, x):
        rms = torch.sqrt(torch.mean(x**2, dim=-1, keepdim=True) + self.eps)
        return x / rms * self.weight

Otra mejora clave es Rotary Position Embeddings (RoPE), que codifica posiciones rotando consultas y claves en el espacio de atención, preservando relaciones relativas mejor que embeddings absolutos. RoPE aplica rotaciones sinusoidales dependientes de la posición.

def apply_rotary_emb(x, freqs):
    x1 = x[..., :x.shape[-1]//2]
    x2 = x[..., x.shape[-1]//2:]
    return torch.cat((x1 * freqs.cos - x2 * freqs.sin, x1 * freqs.sin + x2 * freqs.cos), dim=-1)

# freqs precomputados como cos/sin de ángulos

El KV caching acelera la inferencia auto-regresiva almacenando claves y valores previos, evitando recomputaciones. En código, mantén un buffer que concatena con nuevos KV en cada paso generativo.

Estas enhancements, probadas en modelos como Llama, reducen latencia en un 50% durante despliegues. Integra RoPE en tu atención multi-cabeza modificando las proyecciones Q y K post-lineal.

Adicionalmente, el uso de flash attention optimiza el cálculo de softmax en GPU, fusionando operaciones para menor memoria. Aunque no nativo en PyTorch base, bibliotecas como xformers lo habilitan.

Escalado de Modelos Más Grandes

Escalar un LLM requiere manejar recursos computacionales crecientes. Técnicas como mixed precision training combinan FP16 y FP32 para acelerar cálculos sin perder precisión, usando torch.amp en PyTorch.

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
for batch_x, batch_y in dataloader:
    optimizer.zero_grad()
    with autocast():
        logits = model(batch_x)
        loss = criterion(logits.view(-1, logits.size(-1)), batch_y.view(-1))
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

El logging rico, con herramientas como Weights & Biases, rastrea hiperparámetros y métricas en tiempo real, facilitando debugging.

Para datasets masivos, sharding y distributed data parallel (DDP) distribuyen carga: torch.nn.parallel.DistributedDataParallel(model). En 2025, frameworks como Hugging Face Accelerate simplifican esto para principiantes.

Considera también gradient checkpointing para ahorrar memoria, recomputando activaciones forward durante backward.

Capas de Mixture-of-Experts

Las capas MoE introducen eficiencia al activar sub-redes expertas selectivamente, enroutando tokens a especialistas. Esto escala capacidad sin costo proporcional en parámetros activos.

Una capa MoE básica incluye un router que asigna tokens a top-k expertos vía softmax sobre logits.

class MoELayer(nn.Module):
    def __init__(self, d_model, num_experts, num_experts_per_tok):
        super(MoELayer, self).__init__()
        self.router = nn.Linear(d_model, num_experts)
        self.experts = nn.ModuleList([nn.Linear(d_model, d_model) for _ in range(num_experts)])
        self.num_experts_per_tok = num_experts_per_tok

    def forward(self, x):
        router_logits = self.router(x)
        probs = torch.softmax(router_logits, dim=-1)
        topk_probs, topk_ids = torch.topk(probs, self.num_experts_per_tok, dim=-1)
        topk_probs = topk_probs / topk_probs.sum(dim=-1, keepdim=True)

        output = torch.zeros_like(x)
        for i in range(self.num_experts_per_tok):
            expert_id = topk_ids[..., i]
            prob = topk_probs[..., i].unsqueeze(-1)
            for j, expert in enumerate(self.experts):
                mask = (expert_id == j).unsqueeze(-1)
                output += mask * prob * expert(x)
        return output

MoE reduce el cómputo al procesar solo fracciones del modelo por token, ideal para LLMs de billones de parámetros. Balancea carga con técnicas de auxiliar loss para evitar colapso de routing.

Fine-Tuning Supervisado

El supervised fine-tuning (SFT) adapta el modelo pre-entrenado a tareas específicas mediante datasets curados de pares instrucción-respuesta. Esto refina el comportamiento para utilidad downstream.

Prepara datos como JSONL con campos “prompt” y “completion”. El entrenamiento minimiza cross-entropy sobre completaciones, con prompts como contexto.

class SFTDataset(Dataset):
    def __init__(self, data_file, tokenizer):
        with open(data_file, 'r') as f:
            self.data = [json.loads(line) for line in f]
        self.tokenizer = tokenizer

    def __getitem__(self, idx):
        prompt = self.data[idx]['prompt']
        completion = self.data[idx]['completion']
        input_ids = self.tokenizer.encode(prompt + completion)
        labels = input_ids.copy()
        labels[:len(self.tokenizer.encode(prompt))] = -100  # Ignorar prompt en pérdida
        return torch.tensor(input_ids), torch.tensor(labels)

# Entrenamiento similar al base, pero con labels masked

En sesiones de SFT, usa learning rates bajos (1e-5) y LoRA para eficiencia, adaptando solo adaptadores lineales en lugar del modelo completo.

Esto alinea el modelo hacia outputs deseados, preparando para etapas de alineación más sofisticadas.

Modelado de Recompensas

El reward modeling evalúa la calidad de respuestas generadas, entrenando un clasificador binario o regresión sobre preferencias humanas. Datasets como Anthropic’s HH-RLHF proporcionan comparaciones pairwise.

El modelo de recompensa, típicamente un transformador ligero, predice scores para chosen vs rejected responses. La pérdida usa margin ranking: max(0, score_rejected - score_chosen + margin).

class RewardModel(nn.Module):
    def __init__(self, base_model):
        super(RewardModel, self).__init__()
        self.transformer = base_model
        self.head = nn.Linear(base_model.d_model, 1)

    def forward(self, input_ids):
        hidden = self.transformer(input_ids)[0]  # Última capa
        reward = self.head(hidden.mean(dim=1))  # Pooling mean
        return reward

# Entrenamiento
chosen_rewards = reward_model(chosen_inputs)
rejected_rewards = reward_model(rejected_inputs)
loss = torch.clamp(rejected_rewards - chosen_rewards + margin, min=0).mean()

En 2025, integraciones con APIs de anotación crowdsourced aceleran la recolección de datos. Este paso cuantifica “ayuda” y “inofensividad”, guiando la optimización posterior.

RLHF con PPO

La Reinforcement Learning from Human Feedback (RLHF) refina el modelo usando señales de recompensa vía Proximal Policy Optimization (PPO), un algoritmo de RL estable.

PPO itera actualizando la política (el LLM) para maximizar recompensas esperadas, con KL-divergence penalty para evitar desviaciones extremas de la política inicial.

Implementa un entorno donde el modelo genera respuestas, el reward model las puntúa, y PPO ajusta pesos.

import stable_baselines3 as sb3  # Asumiendo disponibilidad, o implementación custom

# Custom PPO para LLM
class LLMEnv:
    def __init__(self, policy_model, reward_model):
        self.policy = policy_model
        self.reward = reward_model

    def step(self, prompt):
        response = generate_response(self.policy, prompt)  # Función de generación
        reward = self.reward(response)
        return response, reward, done=True, info={}

# Uso de PPO
from sb3 import PPO
env = LLMEnv(model, reward_model)
ppo = PPO("MlpPolicy", env, verbose=1)
ppo.learn(total_timesteps=10000)

PPO clippea actualizaciones de ratio (clipped surrogate objective) para robustez: L*clip = min(r * A, clip(r, 1-eps, 1+eps) _ A), donde r es el ratio de probabilidades nuevas/viejas, A ventaja.

En práctica, entrena en batches de prompts, recolectando trajectories. Esto alinea el modelo hacia outputs preferidas, mitigando toxicidad.

Actualizaciones incluyen value function co-entrenada para estimar baselines, mejorando eficiencia de muestreo.

Conclusiones

En resumen, construir un LLM desde cero abarca desde la arquitectura transformadora hasta técnicas de alineación como RLHF, equipando a desarrolladores con herramientas para innovación en IA. Este recorrido fomenta la experimentación responsable, contribuyendo a un campo en constante evolución. Con las actualizaciones de 2025, enfócate en sostenibilidad y ética para impactos positivos. Explora el código proporcionado para profundizar y adapta a tus proyectos.