Etapa 0: entorno, configuración validada, modelo y entrenador

Base del proyecto ENLACE: un modelo de lenguaje propio entrenado desde cero,
en español, para asistencia general y familiar. El plan completo está en
docs/PLAN.md.

Esta etapa establece el andamiaje y lo verifica de punta a punta:

- Configuración por capas (hardware × model × train × data) validada con
  pydantic. Ningún hiperparámetro vive en el código y una config inválida
  falla al arrancar, no a las tres horas de entrenamiento.
- Perfiles de hardware que aíslan el salto de GPU: la RTX 2060 (Turing) no
  soporta bfloat16 ni FlashAttention-2, así que entrena en float16 con
  GradScaler y backend mem_efficient; el perfil de la 5090 ya está escrito.
  backends.py valida el perfil contra la GPU real antes de empezar.
- Transformer decoder-only estilo Llama: RMSNorm, SwiGLU, RoPE, GQA,
  embeddings atados, QK-norm y z-loss. Los dos últimos son lo que mantiene
  estable el entrenamiento en float16.
- Entrenador con schedule WSD, acumulación de gradiente, precisión mixta,
  checkpointing atómico y reanudación exacta.
- Cargadores de datos con estado serializable: bytes para el smoke test y
  shards uint16 para el corpus real.

48 tests, entre ellos el crítico: reanudar desde un checkpoint reproduce los
pesos de una corrida ininterrumpida, parámetro por parámetro.

Verificado en CPU: 300 pasos sobre texto en español, loss 3.07 -> 1.63.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
2026-07-27 23:04:45 -03:00
commit 90ac2a6582
33 changed files with 2838 additions and 0 deletions
+27
View File
@@ -0,0 +1,27 @@
# RTX 5090 32 GB — Blackwell, compute capability 12.0.
#
# Requiere PyTorch >= 2.7 compilado con CUDA 12.8: las versiones anteriores no
# conocen sm_120 y fallan al compilar kernels. Antes de usar este perfil hay que
# correr el smoke test de la Etapa 0 y verificar el entorno entero.
#
# Con bfloat16 desaparece el GradScaler: bf16 tiene el mismo rango exponente que
# fp32, así que no hay desbordes que escalar.
name: blackwell-5090
device: cuda
dtype: bfloat16
use_grad_scaler: false
attention_backend: flash
compile: true
matmul_precision: high
# 32 GB permiten un micro-lote 4x mayor con la misma acumulación reducida,
# manteniendo un lote efectivo comparable (32 * 4 = 128 secuencias) para que
# las curvas sean comparables entre placas.
micro_batch_size: 32
grad_accum_steps: 4
min_compute_capability: [12, 0]
# Aproximado, y solo para la métrica de MFU: bfloat16 denso (sin sparsity).
# Verificar con un benchmark real al migrar y ajustar este número.
peak_tflops: 209.5
+17
View File
@@ -0,0 +1,17 @@
# CPU — solo para desarrollo y tests. No sirve para entrenar nada real.
#
# Existe para que el código se pueda verificar en la máquina de trabajo, que no
# tiene GPU NVIDIA: los tests de forma, el smoke test a nivel de caracteres y
# la validación de configs corren acá antes de tocar el servidor.
name: cpu
device: cpu
dtype: float32
use_grad_scaler: false
attention_backend: math
compile: false
matmul_precision: high
micro_batch_size: 8
grad_accum_steps: 1
min_compute_capability: null
+28
View File
@@ -0,0 +1,28 @@
# RTX 2060 12 GB — Turing, compute capability 7.5.
#
# Dos límites de la arquitectura Turing determinan este perfil:
#
# 1. No soporta bfloat16. Hay que entrenar en float16, que tiene rango
# exponente reducido y desborda: por eso use_grad_scaler=true. Los picos
# de loss se mitigan además con qk_norm y z_loss en el modelo.
# 2. FlashAttention-2 exige Ampere (sm_80) o superior. El backend
# mem_efficient de PyTorch sí corre acá y usa memoria O(n) igual.
name: turing-2060
device: cuda
dtype: float16
use_grad_scaler: true
attention_backend: mem_efficient
compile: true
matmul_precision: high
# 12 GB con seq_len 2048 y un modelo de ~50M: micro-lote chico y acumulación
# para llegar a un lote efectivo razonable (8 * 16 = 128 secuencias).
micro_batch_size: 8
grad_accum_steps: 16
min_compute_capability: [7, 5]
# Aproximado, y solo para la métrica de MFU: tensor cores en float16 con
# acumulación en float32, que es lo que hace el entrenamiento con AMP.
# Sirve para comparar corridas entre sí, no como dato de hardware.
peak_tflops: 28.0
+23
View File
@@ -0,0 +1,23 @@
# Modelo diminuto a nivel de caracteres, solo para el smoke test de la Etapa 0.
#
# No produce nada útil: existe para validar en ~10 minutos que CUDA, el
# entrenador, el checkpointing, la reanudación y el flujo por SSH funcionan
# antes de comprometer días de cómputo.
name: char-smoke
vocab_size: 256 # se ajusta al vocabulario real del texto al cargar los datos
n_layer: 4
n_head: 4
n_kv_head: 2
d_model: 128
seq_len: 256
ffn_mult: 2.6667
ffn_multiple_of: 32
rope_theta: 10000.0
norm_eps: 1.0e-5
tie_embeddings: true
qk_norm: true
z_loss_weight: 1.0e-4
dropout: 0.0
init_std: 0.02
+30
View File
@@ -0,0 +1,30 @@
# ~50M parámetros. Primer modelo real, pensado para la 2060.
#
# Con vocab 32k y d_model 512, la tabla de embeddings sola son 16.7M
# parámetros: atarla con la capa de salida (tie_embeddings) ahorra un tercio
# del modelo. En modelos chicos esa decisión no es un detalle.
#
# seq_len 2048 y no 1024: el contexto recuperado de las bases del sistema se
# inyecta ahí, así que el presupuesto de contexto es un recurso de primer orden.
name: tiny-50m
vocab_size: 32768
n_layer: 12
n_head: 8
n_kv_head: 2 # GQA 4:1 — reduce la caché KV sin costo medible de calidad
d_model: 512
seq_len: 2048
ffn_mult: 2.6667 # 8/3, el habitual para SwiGLU (tres matrices en vez de dos)
ffn_multiple_of: 64
rope_theta: 10000.0
norm_eps: 1.0e-5
tie_embeddings: true
# qk_norm y z_loss no son opcionales entrenando en float16 en la 2060:
# son lo que evita que el loss explote a mitad de corrida.
qk_norm: true
z_loss_weight: 1.0e-4
dropout: 0.0 # con 3B tokens y 50M params no hay sobreajuste que combatir
init_std: 0.02
+31
View File
@@ -0,0 +1,31 @@
# Pretraining del modelo base (Etapa 2 del plan).
run_name: pretrain-tiny-50m
out_dir: runs
seed: 1337
# ~3B tokens con lote efectivo 128 x 2048 tokens = 262144 tokens por paso.
max_steps: 11500
optimizer:
lr: 6.0e-4
beta1: 0.9
beta2: 0.95
eps: 1.0e-8
weight_decay: 0.1
grad_clip: 1.0
# WarmupStableDecay: el LR sube, se mantiene plano la mayor parte de la
# corrida, y solo decae en el último ~10%. La fase de decay es también donde
# se sube el peso del corpus de estilo (annealing), y es lo que permite
# extender esta corrida más adelante sin recalcular el schedule.
schedule:
kind: wsd
warmup_steps: 500
decay_steps: 1150
min_lr_ratio: 0.0
log_every: 10
eval_every: 500
eval_batches: 20
checkpoint_every: 1000
sample_every: 1000
+29
View File
@@ -0,0 +1,29 @@
# Smoke test de la Etapa 0: minutos, no días.
#
# Criterio de aceptación: el loss baja por debajo de 1.5 y las muestras
# generadas son español reconocible. Si eso pasa, el stack entero está bien.
run_name: smoke-char
out_dir: runs
seed: 1337
max_steps: 2000
optimizer:
lr: 3.0e-3 # modelo diminuto: tolera y necesita un LR alto
beta1: 0.9
beta2: 0.95
eps: 1.0e-8
weight_decay: 0.1
grad_clip: 1.0
schedule:
kind: wsd
warmup_steps: 100
decay_steps: 400
min_lr_ratio: 0.0
log_every: 25
eval_every: 250
eval_batches: 10
checkpoint_every: 500
sample_every: 500