File size: 2,952 Bytes
880f286 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | """Launch entrypoint: build configs + data stream and train.
python run.py --config configs/base_124m.json --data-dir data/fwedu
python run.py --config configs/calibration.json --data-dir data/fwedu
python run.py --config configs/calibration.json --dry-run # synthetic, no data
A config JSON has two objects, "model" and "train", whose keys map directly onto
ModelConfig / TrainConfig fields. CLI --set k=v applies last-mile overrides
(e.g. --set train.batch_size=48) so calibration sweeps don't need new files.
"""
from __future__ import annotations
import sys
import json
import argparse
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent / "src"))
from matilda.config import ModelConfig # noqa: E402
from matilda.train import Trainer, TrainConfig # noqa: E402
from matilda.data import SyntheticStream, BinStream, shard_paths # noqa: E402
def _coerce(value: str):
for cast in (int, float):
try:
return cast(value)
except ValueError:
pass
if value.lower() in ("true", "false"):
return value.lower() == "true"
return value
def apply_overrides(cfg: dict, overrides: list[str]) -> dict:
"""--set train.batch_size=48 -> cfg['train']['batch_size'] = 48"""
for item in overrides:
path, _, raw = item.partition("=")
section, _, key = path.partition(".")
cfg.setdefault(section, {})[key] = _coerce(raw)
return cfg
def build(config: dict):
mcfg = ModelConfig(**config.get("model", {}))
tcfg = TrainConfig(**config.get("train", {}))
return mcfg, tcfg
def build_stream(mcfg, tcfg, data_dir: str | None, dry_run: bool):
if dry_run or not data_dir:
print("[data] synthetic stream (dry run, no real data)")
return SyntheticStream(mcfg.vocab_size, tcfg.batch_size, tcfg.seq_len,
seed=tcfg.seed, device=tcfg.device)
paths = shard_paths(data_dir)
print(f"[data] {len(paths)} shards from {data_dir}")
return BinStream(paths, tcfg.batch_size, tcfg.seq_len, seed=tcfg.seed,
device=tcfg.device)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--config", required=True)
ap.add_argument("--data-dir", default=None)
ap.add_argument("--dry-run", action="store_true")
ap.add_argument("--set", nargs="*", default=[], dest="overrides")
args = ap.parse_args()
config = json.loads(Path(args.config).read_text())
config = apply_overrides(config, args.overrides)
mcfg, tcfg = build(config)
stream = build_stream(mcfg, tcfg, args.data_dir, args.dry_run)
print(f"[run] model={mcfg.d_model}d/{mcfg.n_layers}L "
f"steps={tcfg.total_steps} bs={tcfg.batch_size}x{tcfg.grad_accum} "
f"seq={tcfg.seq_len} compile={tcfg.compile} ckpt={tcfg.ckpt_dir}")
Trainer(mcfg, tcfg, stream).train()
if __name__ == "__main__":
main()
|