| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| """tinynext v2: a small word-piece transformer for in-browser autocomplete that can also fill in the middle. |
| |
| One script does everything: it trains a tokenizer, builds a mix of general text, conversation and email, |
| trains the model, scores it with eval.py (the same harness as tinynext v1 and SmolLM2), exports ONNX with a |
| key/value cache for the browser, and uploads the lot to the model repo under v2/. Training curves, throughput |
| and GPU use go to a Trackio dashboard. |
| |
| hf jobs uv run --namespace huggingface --flavor a100-large --secrets HF_TOKEN train.py |
| hf jobs uv run --namespace huggingface --flavor a10g-small --secrets HF_TOKEN train.py --smoke --no-push |
| |
| The data is prepared on every CPU core before training starts, and the whole token array then sits in GPU |
| memory, so the GPU never waits on Python between steps. |
| """ |
|
|
| import argparse |
| import importlib.util |
| import json |
| import math |
| import os |
| import random |
| import shutil |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| HERE = Path(__file__).parent |
| SPECIAL = ["<|endoftext|>", "<|pre|>", "<|suf|>", "<|mid|>"] |
| EOT, PRE, SUF, MID = range(4) |
|
|
| FINEWEB = "hf://datasets/HuggingFaceFW/fineweb-edu/sample/10BT/{:03d}_00000.parquet" |
| SMOLTALK = "hf://datasets/HuggingFaceTB/smoltalk/data/all/train-{:05d}-of-00009.parquet" |
| ENRON = "hf://datasets/corbt/enron-emails/data/train-{:05d}-of-00003.parquet" |
|
|
|
|
| def cpus() -> int: |
| return len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else os.cpu_count() or 1 |
|
|
|
|
| |
|
|
|
|
| def load_eval_module(repo: str, folder: str): |
| """Uses eval.py and eval_set.json next to this script, or fetches them from the repo, so every model is scored |
| by one implementation on the same text. Registered as a module so worker processes can use it too.""" |
| path = HERE / "eval.py" |
| if not path.exists(): |
| from huggingface_hub import hf_hub_download |
|
|
| for name in ("eval.py", "eval_set.json"): |
| shutil.copy(hf_hub_download(repo, f"{folder}/{name}"), HERE / name) |
| spec = importlib.util.spec_from_file_location("tinynext_eval", path) |
| module = importlib.util.module_from_spec(spec) |
| sys.modules["tinynext_eval"] = module |
| spec.loader.exec_module(module) |
| return module |
|
|
|
|
| |
|
|
|
|
| def to_ascii(text: str) -> str: |
| """The browser model only sees keyboard characters, so it trains on the same.""" |
| import tinynext_eval as ev |
|
|
| return text.translate(ev.SMART_CHARS).encode("ascii", "ignore").decode().strip() |
|
|
|
|
| def smoltalk_text(batch): |
| return {"text": [to_ascii("\n\n".join(m["content"] for m in turns if m["role"] in ("user", "assistant"))) for turns in batch["messages"]]} |
|
|
|
|
| def enron_text(batch): |
| """Training senders only: the eval's held-out senders and enron-autocomplete's validation senders are skipped.""" |
| import hashlib |
|
|
| import tinynext_eval as ev |
|
|
| texts = [] |
| for sender, body in zip(batch["from"], batch["body"]): |
| sender = (sender or "").lower() |
| keep = "@enron.com" in sender and not ev.BULK_SENDER.search(sender) |
| keep = keep and int(hashlib.sha1(sender.encode()).hexdigest(), 16) % 100 >= 10 |
| texts.append(to_ascii("\n\n".join(ev.clean_email(body or ""))) if keep else "") |
| return {"text": texts} |
|
|
|
|
| |
| |
| CHARS_PER_TOKEN = 4.0 |
| TEXT_MARGIN = 1.4 |
|
|
|
|
| def enough_rows(ds, budget_tokens: int, lengths: np.ndarray, seed: int): |
| """A random subset of rows whose text adds up to roughly the token budget, so only that much gets cleaned.""" |
| order = np.random.default_rng(seed).permutation(len(ds)) |
| needed = np.searchsorted(np.cumsum(lengths[order]), budget_tokens * CHARS_PER_TOKEN * TEXT_MARGIN) + 1 |
| return ds.select(np.sort(order[:needed])) |
|
|
|
|
| def load_sources(args, budgets: dict) -> dict: |
| """Downloads parquet files, keeps only as many documents as each source's budget needs, and cleans them on |
| every core.""" |
| import pyarrow.compute as pc |
| from datasets import load_dataset |
|
|
| procs = cpus() |
| fineweb = load_dataset("parquet", data_files=[FINEWEB.format(i) for i in range(args.fineweb_files)], split="train") |
| fineweb = fineweb.select_columns(["text"]) |
| fineweb = enough_rows(fineweb, budgets["fineweb"], pc.utf8_length(fineweb.data.column("text")).to_numpy(), seed=0) |
| fineweb = fineweb.map(lambda b: {"text": [to_ascii(t) for t in b["text"]]}, batched=True, num_proc=procs) |
| smoltalk = load_dataset("parquet", data_files=[SMOLTALK.format(i) for i in range(args.smoltalk_files)], split="train") |
| sample = smoltalk.shuffle(seed=0).select(range(min(2000, len(smoltalk)))) |
| average = np.mean([sum(len(m["content"]) for m in row["messages"]) for row in sample]) |
| smoltalk = enough_rows(smoltalk, budgets["smoltalk"], np.full(len(smoltalk), average), seed=1) |
| smoltalk = smoltalk.map(smoltalk_text, batched=True, num_proc=procs, remove_columns=smoltalk.column_names) |
| enron = load_dataset("parquet", data_files=[ENRON.format(i) for i in range(3)], split="train") |
| enron = enron.map(enron_text, batched=True, num_proc=procs, remove_columns=enron.column_names) |
| sources = {"fineweb": fineweb, "smoltalk": smoltalk, "enron": enron} |
| return {name: ds.filter(lambda b: [len(t) >= 40 for t in b["text"]], batched=True, num_proc=procs) for name, ds in sources.items()} |
|
|
|
|
| def train_tokenizer(sources: dict, weights: dict, vocab: int, sample_docs: int): |
| from tokenizers import Tokenizer, decoders, models, pre_tokenizers, trainers |
|
|
| def sample(): |
| for name, weight in weights.items(): |
| ds = sources[name].shuffle(seed=0) |
| for row in ds.select(range(min(len(ds), int(sample_docs * weight)))): |
| yield row["text"] |
|
|
| tokenizer = Tokenizer(models.BPE()) |
| tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) |
| tokenizer.decoder = decoders.ByteLevel() |
| trainer = trainers.BpeTrainer(vocab_size=vocab, special_tokens=SPECIAL, min_frequency=2, |
| initial_alphabet=pre_tokenizers.ByteLevel.alphabet()) |
| tokenizer.train_from_iterator(sample(), trainer=trainer) |
| return tokenizer |
|
|
|
|
| def fim(ids: list[int], rng: random.Random, max_middle: int, max_suffix: int, max_prefix: int) -> list[int]: |
| """Suffix first, then prefix, then the middle, so typing at the cursor only appends to the prompt.""" |
| if len(ids) < 8: |
| return ids + [EOT] |
| start = rng.randint(1, len(ids) - 2) |
| end = min(len(ids), start + rng.randint(1, max_middle)) |
| prefix, middle, suffix = ids[max(0, start - max_prefix):start], ids[start:end], ids[end:end + max_suffix] |
| if rng.random() < 0.25: |
| suffix = [] |
| return [SUF, *suffix, PRE, *prefix, MID, *middle, EOT] |
|
|
|
|
| def make_encoder(tokenizer_json: str, seq: int, fim_rate: float): |
| """Tokenizes a batch and turns each document into training pieces: plain chunks, or fill-in-the-middle examples.""" |
|
|
| def encode(batch, indices): |
| from tokenizers import Tokenizer |
|
|
| tokenizer = Tokenizer.from_str(tokenizer_json) |
| rng = random.Random(indices[0]) |
| out = [] |
| for encoding in tokenizer.encode_batch(batch["text"]): |
| ids, pieces = encoding.ids, [] |
| for at in range(0, len(ids), seq - 8): |
| chunk = ids[at: at + seq - 8] |
| pieces += fim(chunk, rng, 24, 64, seq - 100) if rng.random() < fim_rate else chunk + [EOT] |
| out.append(pieces) |
| return {"tokens": out} |
|
|
| return encode |
|
|
|
|
| def flatten(ds, budget: int) -> np.ndarray: |
| """All of a source's tokens as one array, cut to its share of the budget (repeated if the source is small).""" |
| import pyarrow.compute as pc |
|
|
| column = ds.data.column("tokens") |
| tokens = pc.list_flatten(column).to_numpy().astype(np.uint16) |
| if len(tokens) < budget: |
| print(f" only {len(tokens) / 1e6:.0f}M tokens, repeating to reach {budget / 1e6:.0f}M", flush=True) |
| tokens = np.tile(tokens, math.ceil(budget / len(tokens))) |
| return tokens[:budget] |
|
|
|
|
| def build_rows(sources: dict, tokenizer, weights: dict, args) -> tuple[np.ndarray, np.ndarray, dict]: |
| encode = make_encoder(tokenizer.to_str(), args.seq, args.fim_rate) |
| train_parts, val_parts, counts = [], [], {} |
| for name, weight in weights.items(): |
| started = time.time() |
| split = sources[name].train_test_split(test_size=0.002, seed=0) |
| encoded = {part: split[part].map(encode, batched=True, batch_size=512, with_indices=True, num_proc=cpus(), |
| remove_columns=["text"]) for part in ("train", "test")} |
| budget = int(args.tokens * weight) |
| train_parts.append(flatten(encoded["train"], budget)) |
| val_parts.append(flatten(encoded["test"], int(args.val_tokens * weight))) |
| counts[name] = {"docs": len(split["train"]), "train_tokens": int(len(train_parts[-1]))} |
| print(f" {name}: {len(split['train']):,} docs → {len(train_parts[-1]) / 1e6:.0f}M tokens in {time.time() - started:.0f}s", flush=True) |
|
|
| def rows(parts): |
| tokens = np.concatenate(parts) |
| count = len(tokens) // (args.seq + 1) |
| return tokens[: count * (args.seq + 1)].reshape(count, args.seq + 1) |
|
|
| return rows(train_parts), rows(val_parts), counts |
|
|
|
|
| |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, dim: int): |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(dim)) |
|
|
| def forward(self, x): |
| return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) * self.weight |
|
|
|
|
| def rotate(x, cos, sin): |
| half = x.shape[-1] // 2 |
| x1, x2 = x[..., :half], x[..., half:] |
| return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, dim: int, heads: int, ffn: int): |
| super().__init__() |
| self.heads, self.head_dim = heads, dim // heads |
| self.norm1, self.norm2 = RMSNorm(dim), RMSNorm(dim) |
| self.qkv = nn.Linear(dim, 3 * dim, bias=False) |
| self.out = nn.Linear(dim, dim, bias=False) |
| self.gate_up = nn.Linear(dim, 2 * ffn, bias=False) |
| self.down = nn.Linear(ffn, dim, bias=False) |
|
|
| def forward(self, x, cos, sin, past_k, past_v, mask): |
| batch, length, _ = x.shape |
| q, k, v = self.qkv(self.norm1(x)).view(batch, length, 3, self.heads, self.head_dim).unbind(2) |
| q, k, v = rotate(q.transpose(1, 2), cos, sin), rotate(k.transpose(1, 2), cos, sin), v.transpose(1, 2) |
| if past_k is not None: |
| k, v = torch.cat([past_k, k], dim=2), torch.cat([past_v, v], dim=2) |
| if mask is None: |
| attended = F.scaled_dot_product_attention(q, k, v, is_causal=True) |
| else: |
| attended = F.scaled_dot_product_attention(q, k, v, attn_mask=mask) |
| x = x + self.out(attended.transpose(1, 2).reshape(batch, length, -1)) |
| gate, up = self.gate_up(self.norm2(x)).chunk(2, dim=-1) |
| return x + self.down(F.silu(gate) * up), k, v |
|
|
|
|
| class TinyNext(nn.Module): |
| def __init__(self, vocab: int, dim: int, layers: int, heads: int, ffn: int): |
| super().__init__() |
| self.config = {"vocab": vocab, "dim": dim, "layers": layers, "heads": heads, "ffn": ffn} |
| self.embed = nn.Embedding(vocab, dim) |
| self.blocks = nn.ModuleList(Block(dim, heads, ffn) for _ in range(layers)) |
| self.norm = RMSNorm(dim) |
| self.head = None |
| self.head_dim = dim // heads |
| inverse = 1.0 / (10000 ** (torch.arange(0, self.head_dim, 2).float() / self.head_dim)) |
| self.register_buffer("inverse_frequency", inverse, persistent=False) |
| nn.init.normal_(self.embed.weight, std=0.02) |
|
|
| def forward(self, ids, past_k=None, past_v=None): |
| """ids: [batch, length]. past_k/past_v: [layers, batch, heads, past, head_dim] or None. |
| Returns logits for every position and the updated cache.""" |
| past = 0 if past_k is None else past_k.shape[3] |
| length = ids.shape[1] |
| positions = torch.arange(length, device=ids.device) + past |
| angles = positions[:, None].float() * self.inverse_frequency[None, :] |
| cos, sin = angles.cos()[None, None], angles.sin()[None, None] |
| mask = None |
| if past_k is not None: |
| keys = torch.arange(past + length, device=ids.device) |
| mask = keys[None, :] <= positions[:, None] |
| x = self.embed(ids) |
| new_k, new_v = [], [] |
| for index, block in enumerate(self.blocks): |
| x, k, v = block(x, cos, sin, None if past_k is None else past_k[index], None if past_v is None else past_v[index], mask) |
| new_k.append(k) |
| new_v.append(v) |
| logits = self.norm(x) @ (self.embed.weight.T if self.head is None else self.head) |
| return logits, torch.stack(new_k), torch.stack(new_v) |
|
|
|
|
| class BrowserStep(nn.Module): |
| """The ONNX graph: always takes a cache (possibly empty) and returns only the last position's logits.""" |
|
|
| def __init__(self, model: TinyNext): |
| super().__init__() |
| self.model = model |
|
|
| def forward(self, ids, past_k, past_v): |
| logits, k, v = self.model(ids, past_k, past_v) |
| return logits[:, -1], k, v |
|
|
|
|
| |
|
|
|
|
| class GpuMonitor: |
| """Reads GPU utilization and memory through NVML, so the logs show whether the GPU is actually busy.""" |
|
|
| def __init__(self): |
| try: |
| import pynvml |
|
|
| pynvml.nvmlInit() |
| self.nvml, self.handle = pynvml, pynvml.nvmlDeviceGetHandleByIndex(0) |
| except Exception as error: |
| print(f"GPU monitoring off: {error}", flush=True) |
| self.nvml = None |
|
|
| def read(self) -> dict: |
| if not self.nvml: |
| return {} |
| rates = self.nvml.nvmlDeviceGetUtilizationRates(self.handle) |
| memory = self.nvml.nvmlDeviceGetMemoryInfo(self.handle) |
| return {"gpu_util_percent": rates.gpu, "gpu_memory_gb": round(memory.used / 1e9, 1)} |
|
|
|
|
| def evaluate_loss(model, rows: torch.Tensor, device, batch: int) -> float: |
| model.eval() |
| losses = [] |
| with torch.no_grad(), torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"): |
| for at in range(0, len(rows), batch): |
| chunk = rows[at: at + batch].long() |
| logits, _, _ = model(chunk[:, :-1]) |
| losses.append(F.cross_entropy(logits.float().flatten(0, 1), chunk[:, 1:].flatten()).item()) |
| model.train() |
| return float(np.mean(losses)) |
|
|
|
|
| def train(model, rows: torch.Tensor, val_rows: torch.Tensor, args, device, tracker) -> list[dict]: |
| steps = int(args.epochs * len(rows) // args.batch) |
| decay = [p for p in model.parameters() if p.dim() >= 2] |
| no_decay = [p for p in model.parameters() if p.dim() < 2] |
| optimizer = torch.optim.AdamW([{"params": decay, "weight_decay": 0.1}, {"params": no_decay, "weight_decay": 0.0}], |
| lr=args.lr, betas=(0.9, 0.95), fused=device.type == "cuda") |
| warmup = max(1, steps // 50) |
| schedule = lambda step: min(1, step / warmup) * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * min(1, step / steps)))) |
| scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, schedule) |
| compiled = torch.compile(model) if device.type == "cuda" else model |
| parameters = sum(p.numel() for p in model.parameters()) |
| gpu, history = GpuMonitor(), [] |
| generator = torch.Generator(device=device).manual_seed(0) |
| window_started, window_tokens, started = time.time(), 0, time.time() |
| print(f"training {steps:,} steps of {args.batch * args.seq:,} tokens", flush=True) |
| for step in range(steps): |
| index = torch.randint(0, len(rows), (args.batch,), device=device, generator=generator) |
| chunk = rows[index].long() |
| with torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"): |
| logits, _, _ = compiled(chunk[:, :-1]) |
| loss = F.cross_entropy(logits.float().flatten(0, 1), chunk[:, 1:].flatten()) |
| optimizer.zero_grad(set_to_none=True) |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) |
| optimizer.step() |
| scheduler.step() |
| window_tokens += args.batch * args.seq |
| if step % args.log_every == 0 or step == steps - 1: |
| loss_value = loss.item() |
| seconds = time.time() - window_started |
| tokens_per_second = window_tokens / max(seconds, 1e-9) |
| entry = { |
| "step": step, "train_loss": round(loss_value, 4), "lr": round(scheduler.get_last_lr()[0], 6), |
| "tokens_per_second": round(tokens_per_second), "minutes": round((time.time() - started) / 60, 2), |
| |
| "mfu_percent": round(100 * 6 * parameters * tokens_per_second / 312e12, 1), **gpu.read(), |
| } |
| if step % (args.log_every * 10) == 0 or step == steps - 1: |
| entry["val_loss"] = round(evaluate_loss(model, val_rows, device, args.batch), 4) |
| history.append(entry) |
| tracker.log(entry, step=step) |
| print(json.dumps(entry), flush=True) |
| window_started, window_tokens = time.time(), 0 |
| return history |
|
|
|
|
| |
|
|
|
|
| def as_token_lm(ev, name: str, model: TinyNext, tokenizer, device, fim: bool = True): |
| texts = [tokenizer.decode([i], skip_special_tokens=False) if i >= len(SPECIAL) else "" for i in range(tokenizer.get_vocab_size())] |
|
|
| @torch.no_grad() |
| def next_logits(ids, cache): |
| past_k, past_v = cache if cache else (None, None) |
| logits, k, v = model(torch.tensor([ids], device=device), past_k, past_v) |
| return logits[0, -1].float().cpu().numpy(), (k, v) |
|
|
| encode = lambda text: tokenizer.encode(text).ids or [EOT] |
| |
| prompt = (lambda prefix, suffix: prefix if not suffix else [SUF, *suffix, PRE, *prefix, MID]) if fim else None |
| return ev.TokenLM(name, next_logits, encode, texts, set(range(len(SPECIAL))), prompt=prompt) |
|
|
|
|
| def parity_prompts(ev, tokenizer, count: int = 300) -> list[list[int]]: |
| """Real text from the eval set, cut at random points, as token prompts.""" |
| rng = random.Random(0) |
| texts = [p for domain in ev.load_eval_set().values() for doc in domain for p in doc] |
| prompts = [] |
| for text in rng.sample(texts, min(count, len(texts))): |
| ids = tokenizer.encode(text[: rng.randint(40, len(text))]).ids[-300:] |
| if len(ids) >= 4: |
| prompts.append(ids) |
| return prompts |
|
|
|
|
| def greedy_tokens(step, ids: list[int], empty, steps: int) -> list[int]: |
| """Prompt, then `steps` cached greedy steps; returns the top token at each point.""" |
| logits, k, v = step(np.array([ids]), empty, empty) |
| tops = [int(logits.argmax())] |
| for _ in range(steps): |
| logits, k, v = step(np.array([[tops[-1]]]), k, v) |
| tops.append(int(logits.argmax())) |
| return tops |
|
|
|
|
| def export(model: TinyNext, output: Path, device, prompts: list[list[int]]) -> dict: |
| """Float and int8 ONNX for the browser, checked against PyTorch on real prompts over several cached steps.""" |
| import onnxruntime as ort |
| from onnxruntime.quantization import QuantType, quantize_dynamic |
|
|
| model = model.float().eval().to("cpu") |
| model.head = nn.Parameter(model.embed.weight.detach().T.contiguous().clone()) |
| cfg = model.config |
| head_dim = cfg["dim"] // cfg["heads"] |
| empty = np.zeros((cfg["layers"], 1, cfg["heads"], 0, head_dim), np.float32) |
| |
| with torch.no_grad(): |
| _, trace_k, trace_v = model(torch.tensor([[1, 2]])) |
| float_path, int8_path = output / "model.onnx", output / "model_int8.onnx" |
| torch.onnx.export( |
| BrowserStep(model), (torch.tensor([[5, 6, 7]]), trace_k, trace_v), float_path, opset_version=17, |
| input_names=["input_ids", "past_k", "past_v"], output_names=["logits", "present_k", "present_v"], |
| dynamic_axes={"input_ids": {1: "length"}, "past_k": {3: "past"}, "past_v": {3: "past"}, |
| "present_k": {3: "total"}, "present_v": {3: "total"}}, |
| dynamo=False, |
| ) |
| quantize_dynamic(float_path, int8_path, weight_type=QuantType.QInt8, op_types_to_quantize=["MatMul", "Gather"], per_channel=True) |
|
|
| @torch.no_grad() |
| def torch_step(ids, k, v): |
| past = None if k.shape[3] == 0 else (torch.from_numpy(k), torch.from_numpy(v)) |
| logits, nk, nv = model(torch.from_numpy(ids), *(past or (None, None))) |
| return logits[0, -1].numpy(), nk.numpy(), nv.numpy() |
|
|
| expected = [greedy_tokens(torch_step, ids, empty, 4) for ids in prompts] |
| report = {"parity_prompts": len(prompts), "parity_steps_per_prompt": 5} |
| for name, path in (("float", float_path), ("int8", int8_path)): |
| session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"]) |
| onnx_step = lambda ids, k, v: (lambda out: (out[0][0], out[1], out[2]))( |
| session.run(None, {"input_ids": ids.astype(np.int64), "past_k": k, "past_v": v})) |
| got = [greedy_tokens(onnx_step, ids, empty, 4) for ids in prompts] |
| first = np.mean([g[0] == e[0] for g, e in zip(got, expected)]) |
| all_steps = np.mean([g == e for g, e in zip(got, expected)]) |
| _, pk, pv = onnx_step(np.array([prompts[0]]), empty, empty) |
| timings = [] |
| for _ in range(200): |
| began = time.perf_counter() |
| onnx_step(np.array([[8]]), pk, pv) |
| timings.append((time.perf_counter() - began) * 1000) |
| report[name] = { |
| "megabytes": round(path.stat().st_size / 1e6, 2), |
| "same_first_token": round(float(first), 4), |
| "same_five_greedy_tokens": round(float(all_steps), 4), |
| "cached_step_ms_p50_cpu": round(float(np.median(timings)), 3), |
| } |
| model.head = None |
| model.to(device) |
| return report |
|
|
|
|
| def headline(scored: dict) -> dict: |
| """The few eval numbers worth a chart: typing saved and 2+ word suggestions at the app's cutoff.""" |
| table = scored["cutoffs"]["0.5"] |
| out = {} |
| for key, m in table.items(): |
| slug = key.replace("/", "_") |
| out[f"eval/{slug}/saved_with_tab"] = m["saved_with_tab"] |
| out[f"eval/{slug}/right"] = m["right"] |
| out[f"eval/{slug}/multiword_shown"] = m["multiword_shown"] |
| return out |
|
|
|
|
| |
|
|
|
|
| def data_key(args) -> str: |
| """Prepared data depends only on these settings, so a training run can find the matching files.""" |
| return f"{args.name}-v{args.vocab}-s{args.seq}-t{int(args.tokens / 1e6)}M-fim{args.fim_rate}-{args.mix}".replace(",", "_").replace("=", "") |
|
|
|
|
| def prepare(args, weights: dict, output: Path): |
| budgets = {name: int((args.tokens + args.val_tokens) * weight) for name, weight in weights.items()} |
| started = time.time() |
| sources = load_sources(args, budgets) |
| print(f"sources ready in {time.time() - started:.0f}s: " + ", ".join(f"{n} {len(d):,} docs" for n, d in sources.items()), flush=True) |
| tokenizer = train_tokenizer(sources, weights, args.vocab, args.tokenizer_docs) |
| tokenizer.save(str(output / "tokenizer.json")) |
| rows, val_rows, counts = build_rows(sources, tokenizer, weights, args) |
| return tokenizer, rows, val_rows, counts |
|
|
|
|
| def upload_prepared(args, output: Path, rows: np.ndarray, val_rows: np.ndarray, counts: dict, minutes: float): |
| from huggingface_hub import HfApi |
|
|
| folder = output / "data" |
| folder.mkdir(exist_ok=True) |
| rows.astype(np.uint16).tofile(folder / "train.bin") |
| val_rows.astype(np.uint16).tofile(folder / "val.bin") |
| shutil.copy(output / "tokenizer.json", folder / "tokenizer.json") |
| meta = {"seq": args.seq, "train_rows": len(rows), "val_rows": len(val_rows), "counts": counts, "minutes": round(minutes, 1), "args": vars(args)} |
| (folder / "meta.json").write_text(json.dumps(meta, indent=1)) |
| api = HfApi(token=os.environ.get("HF_TOKEN")) |
| api.create_repo(args.data_repo, repo_type="dataset", private=True, exist_ok=True) |
| api.upload_folder(folder_path=str(folder), repo_id=args.data_repo, repo_type="dataset", path_in_repo=data_key(args), |
| commit_message=f"prepared data: {data_key(args)}") |
| print(f"Prepared data at https://huggingface.co/datasets/{args.data_repo}/tree/main/{data_key(args)}", flush=True) |
|
|
|
|
| def download_prepared(args, output: Path): |
| from huggingface_hub import snapshot_download |
| from tokenizers import Tokenizer |
|
|
| local = Path(snapshot_download(args.data_repo, repo_type="dataset", allow_patterns=[f"{data_key(args)}/*"], max_workers=16)) |
| folder = local / data_key(args) |
| meta = json.loads((folder / "meta.json").read_text()) |
| shutil.copy(folder / "tokenizer.json", output / "tokenizer.json") |
| rows = np.fromfile(folder / "train.bin", dtype=np.uint16).reshape(meta["train_rows"], meta["seq"] + 1) |
| val_rows = np.fromfile(folder / "val.bin", dtype=np.uint16).reshape(meta["val_rows"], meta["seq"] + 1) |
| return Tokenizer.from_file(str(output / "tokenizer.json")), rows, val_rows, meta["counts"] |
|
|
|
|
| |
|
|
|
|
| def parse_args(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--vocab", type=int, default=8192) |
| parser.add_argument("--dim", type=int, default=320) |
| parser.add_argument("--layers", type=int, default=8) |
| parser.add_argument("--heads", type=int, default=5) |
| parser.add_argument("--ffn", type=int, default=864) |
| parser.add_argument("--seq", type=int, default=512) |
| parser.add_argument("--tokens", type=float, default=1.5e9, help="Training tokens") |
| parser.add_argument("--val-tokens", type=float, default=4e6) |
| parser.add_argument("--epochs", type=float, default=1.0) |
| parser.add_argument("--batch", type=int, default=512) |
| parser.add_argument("--lr", type=float, default=3e-3) |
| parser.add_argument("--fim-rate", type=float, default=0.5) |
| parser.add_argument("--mix", default="fineweb=0.6,smoltalk=0.3,enron=0.1") |
| parser.add_argument("--fineweb-files", type=int, default=2, help="Each is about 2 GB of text, roughly 0.9B tokens") |
| parser.add_argument("--smoltalk-files", type=int, default=9) |
| parser.add_argument("--tokenizer-docs", type=int, default=200_000) |
| parser.add_argument("--eval-pauses", type=int, default=1000) |
| parser.add_argument("--log-every", type=int, default=50) |
| parser.add_argument("--repo", default="cfahlgren1/tinynext") |
| parser.add_argument("--folder", default="v2", help="Where this run's outputs are uploaded") |
| parser.add_argument("--eval-folder", default="v2", help="Where the shared eval.py and eval_set.json live") |
| parser.add_argument("--name", default="tinynext-v2") |
| parser.add_argument("--trackio-space", default="cfahlgren1/tinynext-trackio") |
| parser.add_argument("--stage", choices=("all", "prepare", "train"), default="all", |
| help="prepare: build the tokenizer and token arrays on a CPU job and save them to --data-repo. " |
| "train: download them and train on the GPU. all: both in one job.") |
| parser.add_argument("--data-repo", default="cfahlgren1/tinynext-data") |
| parser.add_argument("--smoke", action="store_true", help="Tiny settings to check the pipeline end to end") |
| parser.add_argument("--no-push", action="store_true") |
| args = parser.parse_args() |
| if args.smoke: |
| args.vocab, args.dim, args.layers, args.heads, args.ffn, args.seq = 2048, 64, 2, 2, 128, 256 |
| args.tokens, args.val_tokens, args.batch, args.tokenizer_docs = 5e6, 2e5, 64, 5000 |
| args.fineweb_files, args.smoltalk_files, args.eval_pauses, args.log_every = 1, 1, 20, 10 |
| args.name = "tinynext-v2-smoke" |
| return args |
|
|
|
|
| class NoTracker: |
| def log(self, *_, **__): |
| pass |
|
|
| def finish(self): |
| pass |
|
|
|
|
| def start_tracker(args): |
| if args.no_push and not os.environ.get("HF_TOKEN"): |
| return NoTracker() |
| import trackio |
|
|
| trackio.init(project="tinynext", name=args.name, space_id=args.trackio_space, config=vars(args)) |
| return trackio |
|
|
|
|
| def main(): |
| args = parse_args() |
| device = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu") |
| torch.manual_seed(0) |
| torch.backends.cuda.matmul.allow_tf32 = True |
| torch.backends.cudnn.allow_tf32 = True |
| output = HERE / "output" |
| output.mkdir(exist_ok=True) |
| ev = load_eval_module(args.repo, args.eval_folder) |
| weights = {k: float(v) for k, v in (part.split("=") for part in args.mix.split(","))} |
| print(f"device {device}, {cpus()} CPUs, mix {weights}", flush=True) |
| tracker = start_tracker(args) |
|
|
| started = time.time() |
| if args.stage == "train": |
| tokenizer, rows, val_rows, counts = download_prepared(args, output) |
| else: |
| tokenizer, rows, val_rows, counts = prepare(args, weights, output) |
| data_minutes = (time.time() - started) / 60 |
| print(f"data ready in {data_minutes:.1f} min: {rows.size / 1e6:.0f}M train tokens, {val_rows.size / 1e6:.1f}M val", flush=True) |
| if args.stage == "prepare": |
| upload_prepared(args, output, rows, val_rows, counts, data_minutes) |
| return |
|
|
| |
| rows_gpu = torch.from_numpy(rows.astype(np.int32)).to(device) |
| val_gpu = torch.from_numpy(val_rows.astype(np.int32)).to(device) |
| model = TinyNext(tokenizer.get_vocab_size(), args.dim, args.layers, args.heads, args.ffn).to(device) |
| parameters = sum(p.numel() for p in model.parameters()) |
| print(f"{parameters / 1e6:.2f}M parameters", flush=True) |
| started = time.time() |
| history = train(model, rows_gpu, val_gpu, args, device, tracker) |
| train_minutes = (time.time() - started) / 60 |
|
|
| model.eval() |
| from safetensors.torch import save_file |
|
|
| save_file({k: v.contiguous() for k, v in model.state_dict().items()}, output / "model.safetensors") |
| scored = ev.score(as_token_lm(ev, args.name, model, tokenizer, device, fim=args.fim_rate > 0), ev.pauses(ev.load_eval_set(), args.eval_pauses)) |
| (output / "metrics").mkdir(exist_ok=True) |
| (output / "metrics" / f"{args.name}.json").write_text(json.dumps(scored, indent=1)) |
| tracker.log(headline(scored)) |
| exported = export(model, output, device, parity_prompts(ev, tokenizer)) |
|
|
| config = { |
| **model.config, "parameters": parameters, "seq": args.seq, "special_tokens": dict(zip(SPECIAL, range(len(SPECIAL)))), |
| "prompt": "No text after the cursor: the text before it. Otherwise: <|suf|> after <|pre|> before <|mid|>.", |
| } |
| (output / "config.json").write_text(json.dumps(config, indent=1)) |
| report = {"args": vars(args), "parameters": parameters, "data": counts, "data_minutes": round(data_minutes, 1), |
| "train_minutes": round(train_minutes, 1), "history": history, "exports": exported} |
| (output / "report.json").write_text(json.dumps(report, indent=1)) |
| print(json.dumps(exported, indent=1)) |
| print(ev.table([output / "metrics" / f"{args.name}.json"])) |
| tracker.finish() |
|
|
| for name in ("train.py", "eval.py", "eval_set.json"): |
| shutil.copy(HERE / name, output / name) |
| if (HERE / "metrics").exists(): |
| for path in (HERE / "metrics").glob("*.json"): |
| if not (output / "metrics" / path.name).exists(): |
| shutil.copy(path, output / "metrics" / path.name) |
| if args.no_push: |
| print(f"Skipping upload; everything is in {output}") |
| return |
| from huggingface_hub import HfApi |
|
|
| HfApi(token=os.environ.get("HF_TOKEN")).upload_folder( |
| folder_path=str(output), repo_id=args.repo, path_in_repo=args.folder, repo_type="model", |
| commit_message=f"{args.name}: {parameters / 1e6:.1f}M params, {train_minutes:.0f} min on GPU", |
| ) |
| print(f"Uploaded to https://huggingface.co/{args.repo}/tree/main/{args.folder}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
