msaldain
90ac2a6582
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>
98 lines
3.2 KiB
Python
98 lines
3.2 KiB
Python
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
|