| """CORTEX training pipeline. |
| |
| Trains a CORTEX model from scratch on code and instruction data. Every dataset must be |
| declared in ``datasets/registry.json`` together with its source and license; the trainer |
| refuses to run on an undeclared or license-incompatible dataset. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import math |
| from dataclasses import dataclass |
| from pathlib import Path |
|
|
| import torch |
|
|
| from model.cortex_model import CortexConfig, count_parameters |
|
|
| __all__ = ["load_registry", "load_jsonl", "CortexTrainer", "train"] |
|
|
| REPO_ROOT = Path(__file__).resolve().parent.parent |
|
|
|
|
| def load_registry(path: str | Path | None = None) -> dict: |
| path = Path(path) if path else REPO_ROOT / "datasets" / "registry.json" |
| return json.loads(path.read_text(encoding="utf-8")) |
|
|
|
|
| def load_jsonl(path: str | Path) -> list[dict]: |
| records = [] |
| with open(path, encoding="utf-8") as fh: |
| for line in fh: |
| line = line.strip() |
| if line: |
| records.append(json.loads(line)) |
| return records |
|
|
|
|
| def assert_dataset_usable(registry: dict, dataset_id: str) -> dict: |
| """Refuse datasets whose license or provenance is not declared. |
| |
| This is the gate that keeps unlicensed or untraceable data out of CORTEX. |
| """ |
| entries = {d["id"]: d for d in registry["datasets"]} |
| if dataset_id not in entries: |
| raise KeyError(f"dataset {dataset_id!r} is not declared in the registry") |
| entry = entries[dataset_id] |
| for field in ("source", "license", "license_url", "origin", "redistribution"): |
| if not entry.get(field): |
| raise ValueError(f"dataset {dataset_id!r} has no {field!r}; refusing to train on it") |
| if entry["license"].lower() in {"unknown", "none", "proprietary"}: |
| raise ValueError( |
| f"dataset {dataset_id!r} has license {entry['license']!r}; refusing to train on it" |
| ) |
| return entry |
|
|
|
|
| def build_token_stream(records: list[dict], tokenizer, seq_len: int): |
| """Turn records into a flat token tensor cut into fixed-length windows.""" |
| ids: list[int] = [] |
| eos = tokenizer.token_to_id("<｜place▁holder▁no▁0｜>") |
| if eos is None: |
| eos = 1 |
| for r in records: |
| text = r.get("text") |
| if text is None: |
| prompt, completion = r.get("prompt", ""), r.get("completion", "") |
| text = f"{prompt}\n{completion}" if (prompt or completion) else "" |
| if not text: |
| continue |
| ids.extend(tokenizer.encode(text, add_special_tokens=False).ids) |
| ids.append(eos) |
| if len(ids) < seq_len + 1: |
| raise ValueError(f"not enough tokens: {len(ids)} for seq_len={seq_len}") |
| n_windows = len(ids) // (seq_len + 1) |
| tensor = torch.tensor(ids[: n_windows * (seq_len + 1)], dtype=torch.long) |
| return tensor.view(n_windows, seq_len + 1) |
|
|
|
|
| @dataclass |
| class TrainState: |
| step: int = 0 |
| tokens_seen: int = 0 |
| best_loss: float = float("inf") |
|
|
|
|
| class CortexTrainer: |
| def __init__(self, model, config: CortexConfig, tokenizer, token_data: torch.Tensor, |
| train_cfg: dict, device: str = "cpu"): |
| self.model = model.to(device) |
| self.config = config |
| self.tokenizer = tokenizer |
| self.data = token_data |
| self.cfg = train_cfg |
| self.device = device |
| self.state = TrainState() |
|
|
| decay, no_decay = [], [] |
| for _, p in model.named_parameters(): |
| if not p.requires_grad: |
| continue |
| (no_decay if p.ndim < 2 else decay).append(p) |
| self.optimizer = torch.optim.AdamW( |
| [ |
| {"params": decay, "weight_decay": train_cfg["weight_decay"]}, |
| {"params": no_decay, "weight_decay": 0.0}, |
| ], |
| lr=train_cfg["learning_rate"], |
| betas=(train_cfg["beta1"], train_cfg["beta2"]), |
| ) |
|
|
| def lr_at(self, step: int) -> float: |
| warmup, total = self.cfg["warmup_steps"], self.cfg["max_steps"] |
| base, floor = self.cfg["learning_rate"], self.cfg["min_learning_rate"] |
| if step < warmup: |
| return base * (step + 1) / max(warmup, 1) |
| progress = (step - warmup) / max(total - warmup, 1) |
| return floor + 0.5 * (base - floor) * (1 + math.cos(math.pi * min(progress, 1.0))) |
|
|
| def train(self, max_steps: int | None = None, on_step=None, on_save=None, |
| save_every: int | None = None, log_every: int | None = None): |
| self.model.train() |
| max_steps = max_steps or self.cfg["max_steps"] |
| save_every = save_every or self.cfg["save_every"] |
| log_every = log_every or self.cfg["log_every"] |
| bs = self.cfg["micro_batch_size"] |
| accum = self.cfg["gradient_accumulation_steps"] |
| seq_len = self.cfg["sequence_length"] |
|
|
| n_windows = self.data.shape[0] |
| g = torch.Generator().manual_seed(self.cfg["seed"]) |
| history = [] |
|
|
| while self.state.step < max_steps: |
| self.optimizer.zero_grad(set_to_none=True) |
| accum_loss = 0.0 |
| for _ in range(accum): |
| idx = torch.randint(0, n_windows, (bs,), generator=g) |
| batch = self.data[idx].to(self.device) |
| inputs, labels = batch[:, :-1], batch[:, 1:].clone() |
| loss = self.model(inputs, labels=labels)["loss"] / accum |
| loss.backward() |
| accum_loss += loss.item() |
| torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.cfg["grad_clip"]) |
| lr = self.lr_at(self.state.step) |
| for group in self.optimizer.param_groups: |
| group["lr"] = lr |
| self.optimizer.step() |
|
|
| self.state.step += 1 |
| self.state.tokens_seen += bs * accum * seq_len |
| self.state.best_loss = min(self.state.best_loss, accum_loss) |
|
|
| if log_every and self.state.step % log_every == 0: |
| history.append({"step": self.state.step, "loss": accum_loss, "lr": lr}) |
| if on_step: |
| on_step(self.state.step, accum_loss, lr) |
| if save_every and self.state.step % save_every == 0 and on_save: |
| on_save(self.state.step, self.state.best_loss) |
|
|
| self.model.eval() |
| return history |
|
|
|
|
| def train(config_path, dataset_ids, tokenizer_path=None, out_dir="checkpoints/dev", |
| max_steps=None, device="cpu", registry_path=None, seed=None, progress=False, |
| save_every=None, log_every=None): |
| """End-to-end training entry point used by the CLI.""" |
| from model.init import init_cortex_model, save_cortex_checkpoint |
| from tokenizer.cortex_tokenizer import load_tokenizer, assert_vocab_matches_config |
|
|
| raw = json.loads(Path(config_path).read_text(encoding="utf-8")) |
| config = CortexConfig.from_json(config_path) |
| config_seed = seed if seed is not None else raw.get("training", {}).get("seed", 1337) |
|
|
| registry = load_registry(registry_path) |
| records = [] |
| for dataset_id in dataset_ids: |
| entry = assert_dataset_usable(registry, dataset_id) |
| ds_path = REPO_ROOT / entry["path"] |
| if not ds_path.exists(): |
| raise FileNotFoundError( |
| f"dataset {dataset_id!r} declared at {entry['path']} is not present. " |
| f"See datasets/README.md for how it is built." |
| ) |
| records.extend(load_jsonl(ds_path)) |
|
|
| tokenizer = load_tokenizer(tokenizer_path) |
| assert_vocab_matches_config(tokenizer, config.vocab_size) |
|
|
| train_cfg = dict(raw["training"]) |
| if max_steps: |
| train_cfg["max_steps"] = max_steps |
| token_data = build_token_stream(records, tokenizer, train_cfg["sequence_length"]) |
|
|
| model = init_cortex_model(config, seed=config_seed) |
| trainer = CortexTrainer(model, config, tokenizer, token_data, train_cfg, device=device) |
|
|
| def dataset_provenance(): |
| return [ |
| {"id": d, "license": assert_dataset_usable(registry, d)["license"], |
| "source": assert_dataset_usable(registry, d)["source"]} |
| for d in dataset_ids |
| ] |
|
|
| def on_step(step, loss, lr): |
| if progress: |
| print(f"step {step:>6} loss {loss:.4f} lr {lr:.2e}", flush=True) |
|
|
| def on_save(step, best): |
| save_cortex_checkpoint( |
| model, config, Path(out_dir) / f"step-{step}", step=step, |
| extra_provenance={"datasets": dataset_provenance(), "training_config": train_cfg}, |
| ) |
| if progress: |
| print(f" saved checkpoint at step {step}", flush=True) |
|
|
| history = trainer.train(max_steps=max_steps, on_step=on_step, on_save=on_save, |
| save_every=save_every, log_every=log_every) |
|
|
| final_dir = save_cortex_checkpoint( |
| model, config, Path(out_dir) / "final", step=trainer.state.step, |
| extra_provenance={ |
| "datasets": dataset_provenance(), |
| "training_config": train_cfg, |
| "tokens_seen": trainer.state.tokens_seen, |
| "final_loss": history[-1]["loss"] if history else None, |
| }, |
| ) |
| return { |
| "checkpoint": str(final_dir), |
| "steps": trainer.state.step, |
| "tokens_seen": trainer.state.tokens_seen, |
| "final_loss": history[-1]["loss"] if history else None, |
| "parameters": count_parameters(model), |
| "history": history, |
| } |
|
|
