| """Pretraining loop: DDP + bf16, WSD schedule, deterministic resume. |
| |
| One run = one (mixture, tokenizer) cell of the design. The mixture is driven |
| entirely by the run config's `langs`/`probs`; the tokenizer by `tok_name`. The |
| loader is globally deterministic and world-size-independent, so a run resumed on |
| a different node count sees the exact same token stream. |
| |
| Cooldown branch: set `branch.from` to a `stable` checkpoint; the schedule then |
| has warmup=stable=0 and only decays, giving a cheap final model at a larger |
| token budget without retraining the trunk. |
| """ |
| import json |
| import os |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
| from .model import ModelConfig, Transformer |
| from .data.loader import MixedStream |
| from .schedule import lr_at, ckpt_interval, stable_end_tokens, total_tokens |
| from .paths import run_dir, ensure |
|
|
|
|
| def _ddp(): |
| if "RANK" in os.environ and int(os.environ.get("WORLD_SIZE", "1")) > 1: |
| import torch.distributed as dist |
| backend = "nccl" if torch.cuda.is_available() else "gloo" |
| dist.init_process_group(backend=backend) |
| rank = dist.get_rank() |
| world = dist.get_world_size() |
| local = int(os.environ.get("LOCAL_RANK", "0")) |
| if torch.cuda.is_available(): |
| torch.cuda.set_device(local) |
| return dist.is_initialized(), rank, world, local |
| return False, 0, 1, 0 |
|
|
|
|
| def _log(rank, path, rec): |
| if rank == 0: |
| with open(path, "a") as f: |
| f.write(json.dumps(rec) + "\n") |
|
|
|
|
| class Trainer: |
| def __init__(self, cfg: dict): |
| self.cfg = cfg |
| self.dist, self.rank, self.world, self.local = _ddp() |
| self.device = torch.device(f"cuda:{self.local}" if torch.cuda.is_available() else "cpu") |
| self.is_cuda = self.device.type == "cuda" |
| torch.manual_seed(cfg.get("seed", 0)) |
| np.random.seed(cfg.get("seed", 0)) |
|
|
| mc = cfg["model"] |
| self.mcfg = ModelConfig(**mc) |
| self.seq_len = self.mcfg.max_seq_len |
| raw_model = Transformer(self.mcfg).to(self.device) |
| if self.rank == 0: |
| print(f"[train] params: {raw_model.num_params(False)/1e6:.1f}M non-embedding, " |
| f"{raw_model.num_params(True)/1e6:.1f}M total") |
| |
| |
| model = raw_model |
| if cfg.get("compile", False) and self.is_cuda: |
| model = torch.compile(model) |
| self.raw_model = raw_model |
| if self.dist: |
| from torch.nn.parallel import DistributedDataParallel as DDP |
| model = DDP(model, device_ids=[self.local] if self.is_cuda else None) |
| self.model = model |
|
|
| opt = cfg.get("optim", {}) |
| self.optim = torch.optim.AdamW( |
| self.raw_model.parameters(), |
| lr=cfg["schedule"]["peak_lr"], |
| betas=tuple(opt.get("betas", (0.9, 0.95))), |
| weight_decay=opt.get("weight_decay", 0.1), |
| eps=opt.get("eps", 1e-8), |
| ) |
| self.grad_clip = opt.get("grad_clip", 1.0) |
|
|
| |
| self.micro_bsz = cfg["train"]["micro_batch_size"] |
| gbt = cfg["train"]["global_batch_tokens"] |
| per_step_windows = max(1, round(gbt / self.seq_len)) |
| |
| unit = self.micro_bsz * self.world |
| self.global_windows = max(unit, (per_step_windows // unit) * unit) |
| self.grad_accum = self.global_windows // unit |
| self.tokens_per_step = self.global_windows * self.seq_len |
| if self.rank == 0: |
| print(f"[train] global batch: {self.global_windows} windows " |
| f"({self.tokens_per_step/1e6:.3f}M tokens), grad_accum={self.grad_accum}") |
|
|
| |
| self.sched = dict(cfg["schedule"]) |
| self.branch = cfg.get("branch") |
| if self.branch: |
| self.sched["warmup_tokens"] = 0 |
| self.sched["stable_tokens"] = 0 |
| self.target_tokens = total_tokens(self.sched) |
| self.ckpt_table = cfg["train"].get("ckpt_schedule", |
| [[2e9, 250e6], [5e9, 500e6], |
| [15e9, 1e9], [1e15, 2e9]]) |
| |
| self.stable_marks = sorted(cfg["train"].get("stable_marks", [])) |
| self.marks_done = set() |
|
|
| |
| self.mixer = MixedStream(cfg["langs"], cfg["tok_name"], self.seq_len, |
| seed=cfg.get("data_seed", 1234), |
| probs=cfg.get("probs")) |
|
|
| self.rdir = ensure(run_dir(cfg["name"])) |
| self.log_path = self.rdir / "train.jsonl" |
| self.tokens = 0 |
| self.step = 0 |
| self.last_ckpt_tokens = 0 |
| self.saved_stable = False |
|
|
| self.wandb = None |
| if self.rank == 0: |
| try: |
| import wandb |
| self.wandb = wandb.init( |
| project="XScript-Pretraining", name=cfg["name"], |
| id=cfg.get("wandb_id", cfg["name"]), |
| resume="allow", config=cfg, |
| ) |
| self.wandb.summary["params_total_M"] = self.raw_model.num_params(True) / 1e6 |
| self.wandb.summary["params_non_embed_M"] = self.raw_model.num_params(False) / 1e6 |
| except Exception as exc: |
| print(f"[train] wandb disabled ({exc})") |
|
|
| |
| def _ckpt_path(self, tag): |
| return ensure(self.rdir / "checkpoints") / f"{tag}.pt" |
|
|
| def save(self, tag, resumable=True): |
| if self.rank != 0: |
| return |
| payload = { |
| "model": self.raw_model.state_dict(), |
| "step": self.step, "tokens": self.tokens, |
| "cfg": self.cfg, |
| } |
| if resumable: |
| payload.update({ |
| "optim": self.optim.state_dict(), |
| "mixer": self.mixer.state_dict(), |
| "last_ckpt_tokens": self.last_ckpt_tokens, |
| "saved_stable": self.saved_stable, |
| "torch_rng": torch.get_rng_state(), |
| }) |
| torch.save(payload, self._ckpt_path(tag)) |
| kind = "full" if resumable else "model-only" |
| print(f"[train] saved {tag} ({kind}) @ {self.tokens/1e9:.3f}B tokens") |
|
|
| def maybe_resume(self): |
| last = self._ckpt_path("last") |
| if self.branch and not last.exists(): |
| ck = torch.load(self.branch["from"], map_location="cpu", weights_only=False) |
| self.raw_model.load_state_dict(ck["model"]) |
| if self.branch.get("load_optim", True): |
| self.optim.load_state_dict(ck["optim"]) |
| if self.rank == 0: |
| print(f"[train] branched from {self.branch['from']} " |
| f"@ {ck['tokens']/1e9:.3f}B (cooldown {self.target_tokens/1e9:.1f}B)") |
| return |
| if last.exists(): |
| ck = torch.load(last, map_location="cpu", weights_only=False) |
| self.raw_model.load_state_dict(ck["model"]) |
| self.optim.load_state_dict(ck["optim"]) |
| self.mixer.load_state_dict(ck["mixer"]) |
| self.step = ck["step"]; self.tokens = ck["tokens"] |
| self.last_ckpt_tokens = ck["last_ckpt_tokens"] |
| self.saved_stable = ck.get("saved_stable", False) |
| self.marks_done = {m for m in self.stable_marks if m <= self.tokens} |
| torch.set_rng_state(ck["torch_rng"]) |
| if self.rank == 0: |
| print(f"[train] resumed @ step {self.step}, {self.tokens/1e9:.3f}B tokens") |
|
|
| |
| def _next_micro_batches(self): |
| """Return grad_accum micro-batches of (x, y) on device for this rank.""" |
| arr, counts = self.mixer.rank_batch(self.global_windows, self.rank, self.world) |
| |
| t = torch.from_numpy(arr.astype(np.int64)) |
| x = t[:, :-1].to(self.device, non_blocking=True) |
| y = t[:, 1:].to(self.device, non_blocking=True) |
| micros = [(x[i:i + self.micro_bsz], y[i:i + self.micro_bsz]) |
| for i in range(0, x.size(0), self.micro_bsz)] |
| return micros, counts |
|
|
| |
| def _eval_sources(self): |
| from .langs import LANGS |
| srcs = {} |
| for l in self.cfg["langs"]: |
| try: |
| from .eval.bpb import load_holdout |
| h = load_holdout(l, self.cfg["train"].get("eval_docs", 500)) |
| if h: |
| srcs[f"holdout_{l}"] = h |
| except Exception: |
| pass |
| try: |
| from . import flores |
| par = flores.load_parallel(list(self.cfg["langs"]), "dev") |
| for l, sents in par.items(): |
| srcs[f"flores_{l}"] = sents |
| except Exception as e: |
| if self.rank == 0: |
| print(f"[train] flores eval skipped: {e}") |
| return srcs |
|
|
| def evaluate(self): |
| if self.rank != 0: |
| return {} |
| from .eval.bpb import eval_sources |
| from .tok.wrapper import Tok |
| from .paths import tokenizer_dir |
| tok = Tok(tokenizer_dir(self.cfg["tok_name"])) |
| res = eval_sources(self.raw_model, tok, self._eval_sources(), |
| self.device, self.seq_len) |
| self.model.train() |
| return res |
|
|
| |
| def train(self): |
| self.maybe_resume() |
| self.model.train() |
| t0 = time.time() |
| log_every = self.cfg["train"].get("log_every", 20) |
| while self.tokens < self.target_tokens: |
| lr = lr_at(self.tokens, self.sched) |
| for g in self.optim.param_groups: |
| g["lr"] = lr |
|
|
| micros, counts = self._next_micro_batches() |
| self.optim.zero_grad(set_to_none=True) |
| loss_acc = 0.0 |
| for j, (x, y) in enumerate(micros): |
| sync = (not self.dist) or (j == len(micros) - 1) |
| ctx = self.model.no_sync() if (self.dist and not sync) else _null() |
| with ctx: |
| with torch.autocast("cuda", dtype=torch.bfloat16) if self.is_cuda else _null(): |
| _, loss = self.model(x, y) |
| (loss / len(micros)).backward() |
| loss_acc += loss.detach().item() / len(micros) |
| torch.nn.utils.clip_grad_norm_(self.raw_model.parameters(), self.grad_clip) |
| self.optim.step() |
|
|
| self.tokens += self.tokens_per_step |
| self.step += 1 |
|
|
| if self.step % log_every == 0: |
| dt = time.time() - t0 |
| tps = self.tokens_per_step * log_every / dt if dt > 0 else 0 |
| rec = { |
| "step": self.step, "tokens": self.tokens, "lr": lr, |
| "loss": loss_acc, "tok_per_s": round(tps), |
| "mix": self.mixer.stats(), |
| } |
| _log(self.rank, self.log_path, rec) |
| if self.wandb: |
| self.wandb.log({**rec, "tokens_b": self.tokens / 1e9}, step=self.step) |
| if self.rank == 0: |
| print(f"[train] step {self.step} | {self.tokens/1e9:.2f}B | " |
| f"loss {loss_acc:.4f} | lr {lr:.2e} | {tps/1e3:.0f}k tok/s") |
| t0 = time.time() |
|
|
| |
| if (not self.branch and not self.saved_stable |
| and self.tokens >= stable_end_tokens(self.sched)): |
| self.save("stable") |
| self.saved_stable = True |
|
|
| |
| for mark in self.stable_marks: |
| if mark not in self.marks_done and self.tokens >= mark: |
| self.save(f"stable_{int(mark/1e6)}M") |
| self.marks_done.add(mark) |
|
|
| |
| if self.tokens - self.last_ckpt_tokens >= ckpt_interval(self.tokens, self.ckpt_table): |
| self.last_ckpt_tokens = self.tokens |
| self.save("last") |
| self.save(f"step{self.step}_{int(self.tokens/1e6)}M", resumable=False) |
| res = self.evaluate() |
| _log(self.rank, self.log_path, |
| {"step": self.step, "tokens": self.tokens, "eval": res}) |
| if self.wandb and res: |
| self.wandb.log({f"eval/{k}_bpb": v["bpb"] for k, v in res.items()} | |
| {f"eval/{k}_ppl": v["ppl_token"] for k, v in res.items()}, |
| step=self.step) |
| if self.rank == 0 and res: |
| brief = {k: round(v["bpb"], 4) for k, v in res.items()} |
| print(f"[eval] {self.tokens/1e9:.2f}B: {brief}") |
|
|
| self.save("last") |
| self.save("final", resumable=False) |
| res = self.evaluate() |
| _log(self.rank, self.log_path, |
| {"step": self.step, "tokens": self.tokens, "eval_final": res}) |
| if self.wandb: |
| if res: |
| self.wandb.log({f"eval_final/{k}_bpb": v["bpb"] for k, v in res.items()} | |
| {f"eval_final/{k}_ppl": v["ppl_token"] for k, v in res.items()}, |
| step=self.step) |
| self.wandb.finish() |
| if self.rank == 0: |
| print(f"[train] DONE {self.cfg['name']} @ {self.tokens/1e9:.2f}B tokens") |
| if self.dist: |
| import torch.distributed as dist |
| dist.destroy_process_group() |
|
|
|
|
| class _null: |
| def __enter__(self): return self |
| def __exit__(self, *a): return False |
|
|
|
|
| def run_from_config(cfg: dict): |
| Trainer(cfg).train() |
|
|
