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:
@@ -0,0 +1,97 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import yaml
|
||||
|
||||
REPO = Path(__file__).resolve().parents[1]
|
||||
|
||||
# Texto en español sintético pero con la estructura correcta (acentos, ñ,
|
||||
# signos de apertura): suficiente para los tests, y sin depender de la red.
|
||||
_MUESTRA = (
|
||||
"ENLACE responde de forma concreta. No adorna. La luz del living está "
|
||||
"encendida y la temperatura del cuarto es de veintiún grados. ¿Querés que "
|
||||
"apague la del pasillo? El calendario tiene una entrada mañana temprano. "
|
||||
"Andrew aprendió despacio, un día a la vez, y así fue construyendo quién "
|
||||
"era. La memoria no está en los pesos: está en la base de datos. "
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def texto_es(tmp_path_factory: pytest.TempPathFactory) -> Path:
|
||||
path = tmp_path_factory.mktemp("datos") / "texto.txt"
|
||||
path.write_text(_MUESTRA * 200, encoding="utf-8")
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def shards_dir(tmp_path_factory: pytest.TempPathFactory) -> Path:
|
||||
"""Dos shards uint16 con tokens predecibles, para verificar el recorrido."""
|
||||
directory = tmp_path_factory.mktemp("shards")
|
||||
total = 0
|
||||
entries = []
|
||||
for i, count in enumerate([4096, 4096]):
|
||||
tokens = np.arange(total, total + count, dtype=np.uint16)
|
||||
name = f"shard-{i:04d}.bin"
|
||||
tokens.tofile(directory / name)
|
||||
entries.append({"file": name, "tokens": count})
|
||||
total += count
|
||||
(directory / "index.json").write_text(
|
||||
json.dumps({"vocab_size": 8192, "shards": entries, "total_tokens": total})
|
||||
)
|
||||
return directory
|
||||
|
||||
|
||||
def write_run_config(
|
||||
tmp_path: Path,
|
||||
texto: Path,
|
||||
*,
|
||||
run_name: str,
|
||||
max_steps: int,
|
||||
checkpoint_every: int,
|
||||
) -> Path:
|
||||
"""Config de corrida mínima en CPU, apuntando a un texto de prueba."""
|
||||
cfg = {
|
||||
"hardware": yaml.safe_load((REPO / "configs/hardware/cpu.yaml").read_text()),
|
||||
"model": yaml.safe_load((REPO / "configs/model/char-smoke.yaml").read_text()),
|
||||
"train": {
|
||||
"run_name": run_name,
|
||||
"out_dir": str(tmp_path / "runs"),
|
||||
"seed": 1337,
|
||||
"max_steps": max_steps,
|
||||
"optimizer": {
|
||||
"lr": 3.0e-3,
|
||||
"beta1": 0.9,
|
||||
"beta2": 0.95,
|
||||
"eps": 1.0e-8,
|
||||
"weight_decay": 0.1,
|
||||
"grad_clip": 1.0,
|
||||
},
|
||||
"schedule": {
|
||||
"kind": "wsd",
|
||||
"warmup_steps": 2,
|
||||
"decay_steps": 2,
|
||||
"min_lr_ratio": 0.0,
|
||||
},
|
||||
"log_every": 1000,
|
||||
"eval_every": 1000,
|
||||
"eval_batches": 2,
|
||||
"checkpoint_every": checkpoint_every,
|
||||
"sample_every": 0,
|
||||
},
|
||||
"data": {
|
||||
"source": "chars",
|
||||
"text_path": str(texto),
|
||||
"val_fraction": 0.05,
|
||||
},
|
||||
}
|
||||
# Modelo aún más chico: los tests tienen que correr en segundos.
|
||||
cfg["model"].update({"n_layer": 2, "d_model": 64, "n_head": 4, "n_kv_head": 2, "seq_len": 64})
|
||||
cfg["hardware"].update({"micro_batch_size": 4, "grad_accum_steps": 2})
|
||||
|
||||
path = tmp_path / f"{run_name}.yaml"
|
||||
path.write_text(yaml.safe_dump(cfg))
|
||||
return path
|
||||
@@ -0,0 +1,126 @@
|
||||
"""La config tiene que fallar al arrancar, no a las tres horas."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
|
||||
from enlace.config.load import ConfigError, compose, load_config
|
||||
|
||||
REPO = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def test_composicion_por_capas():
|
||||
cfg = load_config(REPO / "configs/runs/pretrain-2060.yaml")
|
||||
assert cfg.hardware.name == "turing-2060"
|
||||
assert cfg.model.name == "tiny-50m"
|
||||
assert cfg.train.run_name == "pretrain-tiny-50m"
|
||||
assert cfg.data.source == "shards"
|
||||
|
||||
|
||||
def test_el_yaml_raiz_pisa_las_capas_incluidas():
|
||||
cfg = load_config(REPO / "configs/runs/smoke-2060.yaml")
|
||||
# El perfil de la 2060 declara micro_batch_size 8; la corrida lo sube a 64
|
||||
# porque el modelo char es diminuto.
|
||||
assert cfg.hardware.micro_batch_size == 64
|
||||
assert cfg.hardware.dtype == "float16" # esto sí viene del perfil
|
||||
|
||||
|
||||
def test_overrides_de_linea_de_comandos():
|
||||
cfg = load_config(REPO / "configs/runs/smoke-cpu.yaml", ["train.seed=99"])
|
||||
assert cfg.train.seed == 99
|
||||
|
||||
|
||||
@pytest.mark.parametrize("perfil", ["turing-2060", "blackwell-5090", "cpu"])
|
||||
def test_todos_los_perfiles_de_hardware_son_validos(perfil, tmp_path):
|
||||
"""Un perfil roto solo se descubre al migrar de placa; mejor ahora."""
|
||||
raiz = {
|
||||
"include": {
|
||||
"hardware": f"hardware/{perfil}.yaml",
|
||||
"model": "model/tiny-50m.yaml",
|
||||
"train": "train/pretrain.yaml",
|
||||
"data": "data/corpus.yaml",
|
||||
}
|
||||
}
|
||||
path = REPO / "configs" / "runs" / f"_tmp_{perfil}.yaml"
|
||||
path.write_text(yaml.safe_dump(raiz))
|
||||
try:
|
||||
cfg = load_config(path)
|
||||
assert cfg.hardware.name == perfil
|
||||
finally:
|
||||
path.unlink()
|
||||
|
||||
|
||||
def _raiz(tmp_path: Path, **parches) -> Path:
|
||||
"""Config raíz válida con parches encima, para probar validaciones."""
|
||||
base = {
|
||||
"include": {
|
||||
"hardware": "hardware/cpu.yaml",
|
||||
"model": "model/tiny-50m.yaml",
|
||||
"train": "train/pretrain.yaml",
|
||||
"data": "data/corpus.yaml",
|
||||
}
|
||||
}
|
||||
base.update(parches)
|
||||
path = REPO / "configs" / "runs" / "_tmp_test.yaml"
|
||||
path.write_text(yaml.safe_dump(base))
|
||||
return path
|
||||
|
||||
|
||||
def _espera_error(path: Path, fragmento: str):
|
||||
try:
|
||||
with pytest.raises(ConfigError) as exc:
|
||||
load_config(path)
|
||||
assert fragmento in str(exc.value)
|
||||
finally:
|
||||
path.unlink()
|
||||
|
||||
|
||||
def test_rechaza_campos_desconocidos(tmp_path):
|
||||
# Un typo tiene que ser un error ruidoso, no un default silencioso.
|
||||
_espera_error(_raiz(tmp_path, model={"n_layers": 12}), "n_layers")
|
||||
|
||||
|
||||
def test_rechaza_fp16_sin_grad_scaler(tmp_path):
|
||||
_espera_error(
|
||||
_raiz(tmp_path, hardware={"dtype": "float16", "use_grad_scaler": False}),
|
||||
"use_grad_scaler",
|
||||
)
|
||||
|
||||
|
||||
def test_rechaza_grad_scaler_sin_fp16(tmp_path):
|
||||
_espera_error(
|
||||
_raiz(tmp_path, hardware={"dtype": "float32", "use_grad_scaler": True}),
|
||||
"use_grad_scaler",
|
||||
)
|
||||
|
||||
|
||||
def test_rechaza_dmodel_no_divisible_por_heads(tmp_path):
|
||||
_espera_error(_raiz(tmp_path, model={"d_model": 500}), "no es divisible")
|
||||
|
||||
|
||||
def test_rechaza_gqa_incoherente(tmp_path):
|
||||
_espera_error(_raiz(tmp_path, model={"n_head": 8, "n_kv_head": 3}), "n_kv_head")
|
||||
|
||||
|
||||
def test_rechaza_schedule_que_no_entra(tmp_path):
|
||||
_espera_error(
|
||||
_raiz(tmp_path, train={"max_steps": 100, "schedule": {"warmup_steps": 80, "decay_steps": 80}}),
|
||||
"no quedaría fase estable",
|
||||
)
|
||||
|
||||
|
||||
def test_rechaza_data_source_sin_su_ruta(tmp_path):
|
||||
_espera_error(_raiz(tmp_path, data={"source": "chars", "text_path": None}), "text_path")
|
||||
|
||||
|
||||
def test_include_con_capa_desconocida(tmp_path):
|
||||
path = REPO / "configs" / "runs" / "_tmp_bad.yaml"
|
||||
path.write_text(yaml.safe_dump({"include": {"hardwar": "hardware/cpu.yaml"}}))
|
||||
try:
|
||||
with pytest.raises(ConfigError, match="capas desconocidas"):
|
||||
compose(path)
|
||||
finally:
|
||||
path.unlink()
|
||||
@@ -0,0 +1,143 @@
|
||||
"""Cargadores de datos: determinismo, reanudación y recorrido de shards."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from enlace.data.loaders import ByteStream, ShardStream
|
||||
|
||||
|
||||
def _byte_stream(texto, seed=1337):
|
||||
return ByteStream(
|
||||
text_path=texto,
|
||||
batch_size=2,
|
||||
seq_len=16,
|
||||
val_fraction=0.1,
|
||||
seed=seed,
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
|
||||
|
||||
def test_los_targets_son_la_entrada_desplazada_un_token(texto_es):
|
||||
stream = _byte_stream(texto_es)
|
||||
x, y = stream.next_batch("train")
|
||||
assert x.shape == y.shape == (2, 16)
|
||||
assert torch.equal(x[:, 1:], y[:, :-1])
|
||||
|
||||
|
||||
def test_la_misma_semilla_da_los_mismos_lotes(texto_es):
|
||||
a = _byte_stream(texto_es)
|
||||
b = _byte_stream(texto_es)
|
||||
for _ in range(3):
|
||||
xa, _ = a.next_batch("train")
|
||||
xb, _ = b.next_batch("train")
|
||||
assert torch.equal(xa, xb)
|
||||
|
||||
|
||||
def test_restaurar_el_estado_continua_la_misma_secuencia(texto_es):
|
||||
a = _byte_stream(texto_es)
|
||||
for _ in range(5):
|
||||
a.next_batch("train")
|
||||
estado = a.state_dict()
|
||||
esperado = [a.next_batch("train")[0] for _ in range(3)]
|
||||
|
||||
b = _byte_stream(texto_es)
|
||||
b.load_state_dict(estado)
|
||||
obtenido = [b.next_batch("train")[0] for _ in range(3)]
|
||||
|
||||
for e, o in zip(esperado, obtenido):
|
||||
assert torch.equal(e, o)
|
||||
|
||||
|
||||
def test_evaluar_no_altera_la_secuencia_de_entrenamiento(texto_es):
|
||||
"""Generadores separados por split: cambiar eval_every no debe mover los
|
||||
lotes de entrenamiento, o dos corridas dejarían de ser comparables."""
|
||||
a = _byte_stream(texto_es)
|
||||
esperado = [a.next_batch("train")[0] for _ in range(3)]
|
||||
|
||||
b = _byte_stream(texto_es)
|
||||
lotes = []
|
||||
for _ in range(3):
|
||||
b.next_batch("val")
|
||||
lotes.append(b.next_batch("train")[0])
|
||||
|
||||
for e, o in zip(esperado, lotes):
|
||||
assert torch.equal(e, o)
|
||||
|
||||
|
||||
def test_el_vocabulario_de_bytes_es_siempre_256(texto_es):
|
||||
assert _byte_stream(texto_es).vocab_size == 256
|
||||
|
||||
|
||||
def test_decodifica_utf8_con_acentos():
|
||||
assert ByteStream.decode(list("El niño está acá.".encode())) == "El niño está acá."
|
||||
|
||||
|
||||
def test_texto_inexistente_da_un_error_util(tmp_path):
|
||||
with pytest.raises(FileNotFoundError, match="prepare_smoke_data"):
|
||||
_byte_stream(tmp_path / "no-existe.txt")
|
||||
|
||||
|
||||
def _shard_stream(directory, batch_size=2, seq_len=8):
|
||||
return ShardStream(
|
||||
shards_dir=directory,
|
||||
batch_size=batch_size,
|
||||
seq_len=seq_len,
|
||||
val_fraction=0.1,
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
|
||||
|
||||
def test_los_shards_se_recorren_en_orden(shards_dir):
|
||||
stream = _shard_stream(shards_dir)
|
||||
x, y = stream.next_batch("train")
|
||||
# El fixture escribe tokens consecutivos: 0, 1, 2, ...
|
||||
assert x[0].tolist() == list(range(8))
|
||||
assert y[0].tolist() == list(range(1, 9))
|
||||
assert x[1].tolist() == list(range(8, 16))
|
||||
|
||||
x2, _ = stream.next_batch("train")
|
||||
assert x2[0].tolist() == list(range(16, 24))
|
||||
|
||||
|
||||
def test_la_lectura_cruza_el_limite_entre_shards(shards_dir):
|
||||
"""Un lote que empieza cerca del final de un shard tiene que continuar en el
|
||||
siguiente sin saltarse ni repetir tokens."""
|
||||
stream = _shard_stream(shards_dir, batch_size=1, seq_len=8)
|
||||
stream.load_state_dict({"train": 4090})
|
||||
x, _ = stream.next_batch("train")
|
||||
assert x[0].tolist() == list(range(4090, 4098)) # cruza de shard 0 a shard 1
|
||||
|
||||
|
||||
def test_restaurar_la_posicion_reanuda_donde_iba(shards_dir):
|
||||
a = _shard_stream(shards_dir)
|
||||
for _ in range(4):
|
||||
a.next_batch("train")
|
||||
estado = a.state_dict()
|
||||
esperado, _ = a.next_batch("train")
|
||||
|
||||
b = _shard_stream(shards_dir)
|
||||
b.load_state_dict(estado)
|
||||
obtenido, _ = b.next_batch("train")
|
||||
assert torch.equal(esperado, obtenido)
|
||||
|
||||
|
||||
def test_al_terminar_la_epoca_vuelve_al_principio(shards_dir):
|
||||
stream = _shard_stream(shards_dir, batch_size=1, seq_len=8)
|
||||
primero, _ = stream.next_batch("train")
|
||||
# El split de train son 7372 tokens; posicionarse casi al final.
|
||||
stream.load_state_dict({"train": 7370})
|
||||
stream.next_batch("train")
|
||||
reiniciado, _ = stream.next_batch("train")
|
||||
assert reiniciado[0, 0].item() == 8 # segundo lote de la época nueva
|
||||
|
||||
|
||||
def test_shards_faltantes_dan_un_error_util(tmp_path):
|
||||
with pytest.raises(FileNotFoundError, match="prepare_data"):
|
||||
_shard_stream(tmp_path)
|
||||
|
||||
|
||||
def test_corpus_demasiado_chico_para_el_lote(shards_dir):
|
||||
with pytest.raises(ValueError, match="menos que los"):
|
||||
_shard_stream(shards_dir, batch_size=64, seq_len=1024)
|
||||
@@ -0,0 +1,152 @@
|
||||
"""Propiedades de la arquitectura que tienen que valer en cualquier escala."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from enlace.config.schema import ModelConfig
|
||||
from enlace.model.transformer import Transformer, apply_rope, build_rope_cache
|
||||
|
||||
|
||||
def _cfg(**parches) -> ModelConfig:
|
||||
base = dict(
|
||||
name="test",
|
||||
vocab_size=256,
|
||||
n_layer=2,
|
||||
n_head=4,
|
||||
n_kv_head=2,
|
||||
d_model=64,
|
||||
seq_len=32,
|
||||
)
|
||||
base.update(parches)
|
||||
return ModelConfig(**base)
|
||||
|
||||
|
||||
def test_loss_inicial_es_la_del_azar():
|
||||
"""Un modelo recién inicializado no sabe nada: su loss es ln(vocab).
|
||||
|
||||
Si arranca muy por debajo hay una fuga de información (targets sin
|
||||
desplazar, por ejemplo); si arranca muy por encima, la inicialización está
|
||||
mal escalada y la corrida va a tardar en despegar.
|
||||
"""
|
||||
torch.manual_seed(0)
|
||||
cfg = _cfg()
|
||||
model = Transformer(cfg, "math")
|
||||
x = torch.randint(0, cfg.vocab_size, (8, cfg.seq_len + 1))
|
||||
_, loss = model(x[:, :-1], x[:, 1:])
|
||||
assert loss.item() == pytest.approx(math.log(cfg.vocab_size), abs=0.15)
|
||||
|
||||
|
||||
def test_embeddings_atados_comparten_memoria():
|
||||
model = Transformer(_cfg(tie_embeddings=True), "math")
|
||||
assert model.lm_head.weight.data_ptr() == model.tok_emb.weight.data_ptr()
|
||||
|
||||
suelto = Transformer(_cfg(tie_embeddings=False), "math")
|
||||
assert suelto.lm_head.weight.data_ptr() != suelto.tok_emb.weight.data_ptr()
|
||||
|
||||
|
||||
def test_atar_embeddings_ahorra_una_tabla_entera():
|
||||
atado = Transformer(_cfg(tie_embeddings=True), "math").num_parameters()
|
||||
suelto = Transformer(_cfg(tie_embeddings=False), "math").num_parameters()
|
||||
assert suelto - atado == 256 * 64
|
||||
|
||||
|
||||
def test_la_atencion_es_causal():
|
||||
"""Cambiar un token no puede alterar las predicciones anteriores.
|
||||
|
||||
Es la propiedad que hace que el entrenamiento tenga sentido; si se rompe,
|
||||
el loss baja de forma espectacular y el modelo no sirve para nada.
|
||||
"""
|
||||
torch.manual_seed(0)
|
||||
cfg = _cfg()
|
||||
model = Transformer(cfg, "math").eval()
|
||||
x = torch.randint(0, cfg.vocab_size, (1, cfg.seq_len))
|
||||
with torch.no_grad():
|
||||
base, _ = model(x)
|
||||
alterado = x.clone()
|
||||
alterado[0, -1] = (alterado[0, -1] + 1) % cfg.vocab_size
|
||||
otro, _ = model(alterado)
|
||||
assert torch.allclose(base[:, :-1], otro[:, :-1], atol=1e-6)
|
||||
|
||||
|
||||
def test_ignore_index_excluye_posiciones_enmascaradas():
|
||||
"""El SFT entrena solo sobre los turnos del asistente: el resto va con -100."""
|
||||
torch.manual_seed(0)
|
||||
cfg = _cfg()
|
||||
model = Transformer(cfg, "math")
|
||||
x = torch.randint(0, cfg.vocab_size, (4, cfg.seq_len + 1))
|
||||
inp, tgt = x[:, :-1], x[:, 1:].clone()
|
||||
tgt[:, : cfg.seq_len // 2] = -100
|
||||
_, loss = model(inp, tgt)
|
||||
assert torch.isfinite(loss)
|
||||
|
||||
|
||||
def test_todo_enmascarado_no_rompe():
|
||||
cfg = _cfg()
|
||||
model = Transformer(cfg, "math")
|
||||
x = torch.randint(0, cfg.vocab_size, (2, cfg.seq_len))
|
||||
tgt = torch.full_like(x, -100)
|
||||
_, loss = model(x, tgt)
|
||||
assert torch.isnan(loss) or torch.isfinite(loss) # no debe explotar
|
||||
|
||||
|
||||
def test_rechaza_secuencias_mas_largas_que_el_contexto():
|
||||
cfg = _cfg()
|
||||
model = Transformer(cfg, "math")
|
||||
x = torch.randint(0, cfg.vocab_size, (1, cfg.seq_len + 1))
|
||||
with pytest.raises(ValueError, match="supera seq_len"):
|
||||
model(x)
|
||||
|
||||
|
||||
def test_generate_agrega_exactamente_los_tokens_pedidos():
|
||||
cfg = _cfg()
|
||||
model = Transformer(cfg, "math")
|
||||
inicio = torch.zeros((2, 3), dtype=torch.long)
|
||||
out = model.generate(inicio, max_new_tokens=7, top_k=5)
|
||||
assert out.shape == (2, 10)
|
||||
assert torch.equal(out[:, :3], inicio)
|
||||
|
||||
|
||||
def test_gqa_reduce_las_proyecciones_kv():
|
||||
"""Con GQA 4:1 las matrices de keys/values son un cuarto de las de queries."""
|
||||
cfg = _cfg(n_head=8, n_kv_head=2, d_model=64)
|
||||
model = Transformer(cfg, "math")
|
||||
attn = model.blocks[0].attn
|
||||
assert attn.wq.out_features == 8 * cfg.head_dim
|
||||
assert attn.wk.out_features == 2 * cfg.head_dim
|
||||
assert attn.n_rep == 4
|
||||
|
||||
|
||||
def test_rope_preserva_la_norma():
|
||||
"""RoPE es una rotación: no cambia la magnitud de los vectores."""
|
||||
cos, sin = build_rope_cache(16, 8, 10_000.0)
|
||||
x = torch.randn(2, 3, 16, 8)
|
||||
y = apply_rope(x, cos, sin)
|
||||
assert torch.allclose(x.norm(dim=-1), y.norm(dim=-1), atol=1e-5)
|
||||
|
||||
|
||||
def test_z_loss_penaliza_logits_grandes():
|
||||
torch.manual_seed(0)
|
||||
cfg_con = _cfg(z_loss_weight=1.0)
|
||||
cfg_sin = _cfg(z_loss_weight=0.0)
|
||||
torch.manual_seed(0)
|
||||
con = Transformer(cfg_con, "math")
|
||||
torch.manual_seed(0)
|
||||
sin = Transformer(cfg_sin, "math")
|
||||
x = torch.randint(0, 256, (4, 33))
|
||||
_, loss_con = con(x[:, :-1], x[:, 1:])
|
||||
_, loss_sin = sin(x[:, :-1], x[:, 1:])
|
||||
assert loss_con.item() > loss_sin.item()
|
||||
|
||||
|
||||
def test_el_modelo_del_plan_pesa_lo_esperado():
|
||||
"""tiny-50m tiene que estar cerca de 50M: es lo que entra en la 2060."""
|
||||
from enlace.config.load import load_config
|
||||
from pathlib import Path
|
||||
|
||||
cfg = load_config(Path(__file__).resolve().parents[1] / "configs/runs/pretrain-2060.yaml")
|
||||
model = Transformer(cfg.model, "math")
|
||||
assert 45e6 < model.num_parameters() < 55e6
|
||||
@@ -0,0 +1,80 @@
|
||||
"""La prueba más importante del entrenamiento.
|
||||
|
||||
Reanudar desde un checkpoint tiene que producir exactamente los mismos pesos que
|
||||
una corrida ininterrumpida. Si no, el daño es silencioso: las curvas se ven
|
||||
bien, el modelo entrena de más sobre datos ya vistos, y nadie se entera hasta
|
||||
que la corrida de dos días terminó y el resultado no sirve.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
from enlace.config.load import load_config
|
||||
from enlace.train import checkpoint
|
||||
from enlace.train.train import train
|
||||
from tests.conftest import write_run_config
|
||||
|
||||
|
||||
def _final_weights(run_dir, step: int):
|
||||
payload = torch.load(run_dir / f"ckpt-{step:08d}.pt", map_location="cpu", weights_only=False)
|
||||
return payload["model"], payload
|
||||
|
||||
|
||||
def test_reanudar_reproduce_la_corrida_completa(tmp_path, texto_es):
|
||||
# Corrida A: 20 pasos sin interrupción.
|
||||
cfg_a = load_config(
|
||||
write_run_config(tmp_path, texto_es, run_name="A", max_steps=20, checkpoint_every=10)
|
||||
)
|
||||
dir_a = train(cfg_a)
|
||||
|
||||
# Corrida B: la misma config, simulando que el proceso murió en el paso 10.
|
||||
# Se copia el checkpoint del paso 10 y se reanuda desde ahí.
|
||||
#
|
||||
# La config tiene que ser idéntica, incluido max_steps: el schedule WSD
|
||||
# ubica el inicio del decay en max_steps - decay_steps, así que reanudar
|
||||
# con otro max_steps cambia el LR de los pasos que faltan. No es un defecto
|
||||
# del checkpointing, pero sí una trampa fácil de pisar.
|
||||
cfg_b = load_config(
|
||||
write_run_config(tmp_path, texto_es, run_name="B", max_steps=20, checkpoint_every=10)
|
||||
)
|
||||
dir_b = Path(cfg_b.train.out_dir) / cfg_b.train.run_name
|
||||
dir_b.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copy(dir_a / "ckpt-00000010.pt", dir_b / "ckpt-00000010.pt")
|
||||
train(cfg_b, resume=True)
|
||||
|
||||
pesos_a, payload_a = _final_weights(dir_a, 20)
|
||||
pesos_b, payload_b = _final_weights(dir_b, 20)
|
||||
|
||||
assert payload_a["step"] == payload_b["step"] == 20
|
||||
assert set(pesos_a) == set(pesos_b)
|
||||
for nombre, tensor_a in pesos_a.items():
|
||||
assert torch.equal(tensor_a, pesos_b[nombre]), (
|
||||
f"el parámetro {nombre} difiere tras reanudar: la reanudación no es exacta"
|
||||
)
|
||||
|
||||
|
||||
def test_el_checkpoint_guarda_la_posicion_del_stream(tmp_path, texto_es):
|
||||
cfg = load_config(
|
||||
write_run_config(tmp_path, texto_es, run_name="C", max_steps=4, checkpoint_every=4)
|
||||
)
|
||||
run_dir = train(cfg)
|
||||
payload = torch.load(run_dir / "ckpt-00000004.pt", map_location="cpu", weights_only=False)
|
||||
|
||||
# Sin el estado del stream, reanudar volvería al principio del corpus.
|
||||
assert "stream" in payload and payload["stream"], "el checkpoint no guardó el stream"
|
||||
assert "train" in payload["stream"]
|
||||
# Y sin la config, el checkpoint no sería reproducible.
|
||||
assert payload["config"]["model"]["name"] == "char-smoke"
|
||||
|
||||
|
||||
def test_checkpoint_atomico_no_deja_archivos_truncados(tmp_path, texto_es):
|
||||
cfg = load_config(
|
||||
write_run_config(tmp_path, texto_es, run_name="D", max_steps=6, checkpoint_every=6)
|
||||
)
|
||||
run_dir = train(cfg)
|
||||
assert checkpoint.latest(run_dir) is not None
|
||||
assert not list(run_dir.glob("*.tmp")), "quedaron temporales de escritura"
|
||||
@@ -0,0 +1,56 @@
|
||||
"""Forma del schedule WSD."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from enlace.config.schema import ScheduleConfig
|
||||
from enlace.train.schedules import build_lr_fn
|
||||
|
||||
|
||||
def _lr_fn(max_steps=1000, warmup=100, decay=200, min_ratio=0.0, base=1e-3):
|
||||
cfg = ScheduleConfig(
|
||||
kind="wsd", warmup_steps=warmup, decay_steps=decay, min_lr_ratio=min_ratio
|
||||
)
|
||||
return build_lr_fn(cfg, base, max_steps)
|
||||
|
||||
|
||||
def test_warmup_sube_linealmente_desde_arriba_de_cero():
|
||||
lr = _lr_fn()
|
||||
assert lr(0) > 0.0 # el primer paso no se desperdicia con lr=0
|
||||
assert lr(0) < lr(50) < lr(99)
|
||||
assert lr(99) == 1e-3
|
||||
|
||||
|
||||
def test_la_fase_estable_es_plana():
|
||||
lr = _lr_fn()
|
||||
assert lr(100) == lr(500) == lr(799) == 1e-3
|
||||
|
||||
|
||||
def test_el_decay_baja_de_forma_monotona_hasta_cero():
|
||||
lr = _lr_fn()
|
||||
valores = [lr(s) for s in range(800, 1000)]
|
||||
assert all(a >= b for a, b in zip(valores, valores[1:]))
|
||||
assert valores[0] == 1e-3
|
||||
assert valores[-1] < 1e-4
|
||||
|
||||
|
||||
def test_min_lr_ratio_pone_un_piso():
|
||||
lr = _lr_fn(min_ratio=0.1)
|
||||
assert lr(999) >= 1e-4 * 0.99
|
||||
|
||||
|
||||
def test_sin_decay_el_lr_queda_plano_hasta_el_final():
|
||||
lr = _lr_fn(decay=0)
|
||||
assert lr(999) == 1e-3
|
||||
|
||||
|
||||
def test_extender_la_corrida_mueve_el_inicio_del_decay():
|
||||
"""Documenta la trampa: max_steps define dónde empieza a decaer el LR.
|
||||
|
||||
Reanudar una corrida interrumpida exige el mismo max_steps; extenderla es
|
||||
una decisión distinta, que hay que tomar ramificando desde un checkpoint
|
||||
anterior al decay.
|
||||
"""
|
||||
corto = _lr_fn(max_steps=1000)
|
||||
largo = _lr_fn(max_steps=2000)
|
||||
assert corto(850) < corto(700) # ya está decayendo
|
||||
assert largo(850) == largo(700) # todavía en la fase estable
|
||||
Reference in New Issue
Block a user