tinynext-v2-3b: 12.5M params, 56 min on GPU

v2/experiments/tinynext-v2-3b/config.json

@@ -0,0 +1,16 @@
+{
+ "vocab": 8192,
+ "dim": 320,
+ "layers": 8,
+ "heads": 5,
+ "ffn": 864,
+ "parameters": 12539200,
+ "seq": 512,
+ "special_tokens": {
+  "<|endoftext|>": 0,
+  "<|pre|>": 1,
+  "<|suf|>": 2,
+  "<|mid|>": 3
+ },
+ "prompt": "No text after the cursor: the text before it. Otherwise: <|suf|> after <|pre|> before <|mid|>."
+}

v2/experiments/tinynext-v2-3b/eval.py

@@ -0,0 +1,499 @@
+# /// script
+# requires-python = ">=3.11,<3.14"
+# dependencies = [
+#   "torch>=2.6,<2.8",
+#   "datasets>=3.0",
+#   "numpy<3",
+#   "onnxruntime>=1.20",
+#   "transformers>=4.46",
+#   "tokenizers>=0.20",
+# ]
+# ///
+"""Scores autocomplete models the way people use them: type real text, pause, and see if the gray suggestion is right.
+
+Every model sees the same pauses in the same text and is decoded with the same rule as offline.page:
+follow the model's top pick while the running product of probabilities stays above a cutoff, and show
+the text up to the last finished word. A suggestion is right when the text really continues that way.
+
+Two scenarios:
+  typing   the cursor is at the end of what's written so far.
+  editing  the cursor is mid-sentence and real text follows a gap of 2-6 words. Only fill-in-the-middle
+           models are given that text; the rest see only what's before the cursor.
+
+Usage:
+  uv run eval.py --model tinynext-v1 --onnx model_quint8.onnx
+  uv run eval.py --model smollm2-135m --hf HuggingFaceTB/SmolLM2-135M
+  uv run eval.py --table                  # compare every metrics/*.json
+"""
+
+import argparse
+import hashlib
+import json
+import random
+import re
+import statistics
+import time
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import numpy as np
+
+HERE = Path(__file__).parent
+EVAL_SET = HERE / "eval_set.json"
+METRICS = HERE / "metrics"
+
+CUTOFFS = (0.3, 0.4, 0.5, 0.6, 0.7)
+MAIN_CUTOFF = 0.5  # what offline.page uses
+PAUSE_EVERY = 12  # characters between pauses
+MIN_SHOWN_CHARS = 4  # offline.page hides shorter suggestions
+MAX_WORDS = 6
+MAX_CHARS = 48
+DOMAINS = ("articles", "messages", "email")
+SCENARIOS = ("typing", "editing")
+PUNCTUATION = ".,!?;:"
+
+# ---- Eval text -------------------------------------------------------------------------------------------------
+# Kept apart from training: Wikipedia isn't in either model's data, OpenAssistant validation isn't in SmolTalk's
+# training split, and the emails come from the senders the original enron-autocomplete held out for its test set.
+
+CUT_REPLY = re.compile(
+    r"(?im)^\s*(?:-+\s*original message|-+\s*forwarded by|-+\s*inline attachment|from:\s|to:\s|sent:\s"
+    r"|subject:\s|cc:\s|.*\bwrote:\s*$|_{5,}|-{5,}|={5,}|\*{5,}|\d{1,2}/\d{1,2}/\d{2,4}\s+\d{1,2}:\d{2})"
+)
+SENSITIVE = re.compile(r"(?i)https?://|www\.|\S+@\S+|\(?\b\d{3}\)?[-. ]\d{3}[-. ]\d{4}\b|\b\d{3}-\d{2}-\d{4}\b")
+BULK_SENDER = re.compile(r"(?i)announce|newsletter|mailer|outlook\.team|^all[._]|general\.|chairman|enron\.(?:on|net)")
+SMART_CHARS = str.maketrans({"‘": "'", "’": "'", "“": '"', "”": '"', "–": "-", "—": "-", "\t": " "})
+
+
+def clean_email(body: str) -> list[str]:
+    """The same cleaning enron-autocomplete was trained with."""
+    text = body.replace("\r\n", "\n").replace("=20", " ").replace("=\n", "").translate(SMART_CHARS)
+    text = "\n".join(line for line in text.split("\n") if not line.lstrip().startswith(">"))
+    cut = CUT_REPLY.search(text)
+    if cut:
+        text = text[: cut.start()]
+    text = re.sub(r"(?m)^-- *\n[\s\S]*", "", text)
+    paragraphs = []
+    for block in re.split(r"\n\s*\n", text):
+        lines = [re.sub(r" +", " ", line).strip() for line in block.split("\n") if line.strip()]
+        if not lines:
+            continue
+        joined = lines[0]
+        for previous, line in zip(lines, lines[1:]):
+            joined += (" " if len(previous) >= 50 else "\n") + line
+        if SENSITIVE.search(joined) or sum(c.isalpha() for c in joined) < 0.6 * len(joined):
+            continue
+        paragraphs.append(joined)
+    return paragraphs
+
+
+def is_held_out_sender(sender: str) -> bool:
+    """enron-autocomplete put senders whose hash bucket is under 5 in its test split."""
+    sender = sender.lower()
+    if "@enron.com" not in sender or BULK_SENDER.search(sender):
+        return False
+    return int(hashlib.sha1(sender.encode()).hexdigest(), 16) % 100 < 5
+
+
+def ascii_paragraphs(text: str, min_chars: int = 40) -> list[str]:
+    text = text.translate(SMART_CHARS).encode("ascii", "ignore").decode()
+    return [p.strip() for p in re.split(r"\n\s*\n|\n", text) if len(p.strip()) >= min_chars]
+
+
+def build_eval_set(docs_per_domain: int = 80) -> dict[str, list[list[str]]]:
+    from datasets import load_dataset
+
+    sets: dict[str, list[list[str]]] = {domain: [] for domain in DOMAINS}
+    wiki = load_dataset(
+        "wikimedia/wikipedia", data_files="20231101.en/train-00040-of-00041.parquet", split="train", streaming=True
+    )
+    for row in wiki:
+        paragraphs = [p for p in ascii_paragraphs(row["text"], 80) if not p.endswith(":")][:6]
+        if len(paragraphs) >= 3:
+            sets["articles"].append(paragraphs)
+        if len(sets["articles"]) >= docs_per_domain:
+            break
+
+    for row in load_dataset("OpenAssistant/oasst1", split="validation"):
+        if row["lang"] == "en" and "```" not in row["text"]:
+            paragraphs = ascii_paragraphs(row["text"])[:6]
+            if paragraphs:
+                sets["messages"].append(paragraphs)
+        if len(sets["messages"]) >= docs_per_domain:
+            break
+
+    for row in load_dataset("corbt/enron-emails", split="train", streaming=True):
+        if is_held_out_sender(row["from"] or ""):
+            paragraphs = [p for p in clean_email(row["body"] or "") if len(p) >= 40][:6]
+            if paragraphs:
+                sets["email"].append(paragraphs)
+        if len(sets["email"]) >= docs_per_domain:
+            break
+    return sets
+
+
+def load_eval_set() -> dict[str, list[list[str]]]:
+    if not EVAL_SET.exists():
+        EVAL_SET.write_text(json.dumps(build_eval_set(), indent=1))
+    return json.loads(EVAL_SET.read_text())
+
+
+@dataclass
+class Pause:
+    domain: str
+    scenario: str
+    before: str  # the document up to the cursor
+    after: str  # text after the cursor that a model may see (empty when typing)
+    truth: str  # what the person actually typed next
+
+
+def pauses(eval_set: dict[str, list[list[str]]], max_per_domain: int) -> list[Pause]:
+    rng = random.Random(0)
+    found: list[Pause] = []
+    for domain in DOMAINS:
+        in_domain: list[Pause] = []
+        for paragraphs in eval_set[domain]:
+            for index, paragraph in enumerate(paragraphs):
+                earlier = "\n".join(paragraphs[:index])
+                for at in range(PAUSE_EVERY, len(paragraph) - 8, PAUSE_EVERY):
+                    before = (earlier + "\n" if earlier else "") + paragraph[:at]
+                    truth = paragraph[at:]
+                    in_domain.append(Pause(domain, "typing", before, "", truth))
+                    # Editing: a gap of 2-6 words, then the rest of the paragraph as it was written.
+                    gap_words = rng.randint(2, 6)
+                    gap = re.match(r"(\S*\s+){%d}" % gap_words, truth)
+                    if gap and len(truth) > gap.end() + 10:
+                        in_domain.append(Pause(domain, "editing", before, truth[gap.end():], truth))
+        rng.shuffle(in_domain)
+        found += [p for p in in_domain if p.scenario == "typing"][:max_per_domain]
+        found += [p for p in in_domain if p.scenario == "editing"][:max_per_domain]
+    return found
+
+
+# ---- Decoding -------------------------------------------------------------------------------------------------
+
+
+@dataclass
+class Trajectory:
+    """Suggestions that end at a word boundary, with the model's confidence in each. Greedy decoding gives the
+    same path for every cutoff, so one decode answers all of them."""
+
+    candidates: list[tuple[str, float]] = field(default_factory=list)
+    milliseconds: float = 0.0
+
+    def at(self, cutoff: float) -> str:
+        best = ""
+        for text, confidence in self.candidates:
+            if confidence >= cutoff:
+                best = text
+        return best.rstrip()
+
+
+class Model:
+    name: str
+    sees_after: bool = False
+
+    def suggest(self, before: str, after: str) -> Trajectory:
+        raise NotImplementedError
+
+
+class CharGRU(Model):
+    """enron-autocomplete / tinynext v1: one character per step, reading only the current paragraph."""
+
+    HIDDEN, END, NEWLINE, FIRST, OFFSET = 512, 1, 2, 3, 29
+
+    def __init__(self, name: str, onnx_path: str):
+        import onnxruntime as ort
+
+        self.name = name
+        self.session = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
+        self.read_text: str | None = None
+        self.reading = None
+
+    def step(self, char_id: int, state):
+        logits, state = self.session.run(None, {"input_id": np.array([char_id], np.int64), "state": state})
+        return logits[0], state
+
+    def encode(self, char: str) -> int | None:
+        if char == "\n":
+            return self.NEWLINE
+        code = ord(char)
+        return code - self.OFFSET if 32 <= code <= 126 else None
+
+    def read(self, text: str):
+        if self.read_text is None or not text.startswith(self.read_text):
+            self.reading = self.step(self.END, np.zeros((1, self.HIDDEN), np.float32))
+            self.read_text = ""
+        for char in text[len(self.read_text):]:
+            char_id = self.encode(char)
+            if char_id is not None:
+                self.reading = self.step(char_id, self.reading[1])
+        self.read_text = text
+        return self.reading
+
+    def suggest(self, before: str, after: str) -> Trajectory:
+        began = time.perf_counter()
+        logits, state = self.read(before.split("\n")[-1])
+        trajectory, run, confidence, words = Trajectory(), "", 1.0, 0
+        while len(run) < MAX_CHARS and confidence >= min(CUTOFFS):
+            probabilities = np.exp(logits - logits.max())
+            probabilities /= probabilities.sum()
+            char_id = int(probabilities.argmax())
+            if char_id < self.FIRST:
+                if run:
+                    trajectory.candidates.append((run, confidence))
+                break
+            char = chr(char_id + self.OFFSET)
+            if char == " ":
+                trajectory.candidates.append((run, confidence))
+                words += 1
+                if words >= MAX_WORDS:
+                    break
+            confidence *= float(probabilities[char_id])
+            run += char
+            if char in PUNCTUATION:
+                trajectory.candidates.append((run, confidence))
+            logits, state = self.step(char_id, state)
+        trajectory.milliseconds = (time.perf_counter() - began) * 1000
+        return trajectory
+
+
+class TokenLM(Model):
+    """A word-piece model. `next_logits(ids, cache)` runs one forward pass and returns (last logits, cache).
+
+    The word being typed is "healed": the prompt stops before it, and the first generated pieces must spell it
+    out, so a half-typed word like "resig" can still become "resignation" in one piece.
+    """
+
+    def __init__(self, name: str, next_logits: Callable, encode: Callable[[str], list[int]], token_texts: list[str],
+                 stop_ids: set[int], prompt: Callable[[list[int], list[int]], list[int]] | None = None,
+                 context_chars: int = 1500, suffix_chars: int = 300):
+        self.name = name
+        self.next_logits = next_logits
+        self.encode = encode
+        self.token_texts = token_texts
+        self.stop_ids = stop_ids
+        self.prompt = prompt
+        self.sees_after = prompt is not None
+        self.context_chars = context_chars
+        self.suffix_chars = suffix_chars
+        lengths = np.array([len(t) for t in token_texts])
+        self.nonempty = lengths > 0
+
+    def allowed(self, need: str) -> np.ndarray:
+        """Pieces consistent with the rest of the half-typed word."""
+        return np.array([bool(t) and (need.startswith(t) or t.startswith(need)) for t in self.token_texts])
+
+    def suggest(self, before: str, after: str) -> Trajectory:
+        began = time.perf_counter()
+        before = before[-self.context_chars:]
+        # The half-typed word keeps its leading space, since word pieces start with one. A line break stays in the prompt.
+        space, newline = before.rfind(" "), before.rfind("\n")
+        if space > newline:
+            stem, need = before[:space], before[space:]
+        else:
+            stem, need = before[: newline + 1], before[newline + 1:]
+        prefix_ids = self.encode(stem)
+        if self.prompt:
+            ids = self.prompt(prefix_ids, self.encode(after[: self.suffix_chars]) if after else [])
+        else:
+            ids = prefix_ids
+        logits, cache = self.next_logits(ids, None)
+
+        trajectory, generated, confidence = Trajectory(), "", 1.0
+        for _ in range(16):
+            logits = logits.astype(np.float64)
+            if need:
+                mask = self.allowed(need)
+                if not mask.any():
+                    break
+                logits = np.where(mask, logits, -np.inf)
+            probabilities = np.exp(logits - logits.max())
+            probabilities /= probabilities.sum()
+            token = int(probabilities.argmax())
+            text = self.token_texts[token]
+            if token in self.stop_ids or "\n" in text:
+                if generated:
+                    trajectory.candidates.append((generated, confidence))
+                break
+            if need:
+                consumed = min(len(need), len(text))
+                text, need = text[consumed:], need[consumed:]
+            if text[:1].isspace() and generated:
+                trajectory.candidates.append((generated, confidence))
+            confidence *= float(probabilities[token])
+            generated += text
+            if generated.rstrip()[-1:] in PUNCTUATION and generated.strip():
+                trajectory.candidates.append((generated, confidence))
+            if confidence < min(CUTOFFS) or len(generated) > MAX_CHARS or len(generated.split()) > MAX_WORDS:
+                break
+            logits, cache = self.next_logits([token], cache)
+        trajectory.milliseconds = (time.perf_counter() - began) * 1000
+        return trajectory
+
+
+def hf_causal_lm(name: str, repo: str) -> TokenLM:
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    device = "mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu"
+    tokenizer = AutoTokenizer.from_pretrained(repo)
+    model = AutoModelForCausalLM.from_pretrained(repo, torch_dtype=torch.float32).to(device).eval()
+
+    @torch.no_grad()
+    def next_logits(ids, cache):
+        out = model(torch.tensor([ids], device=device), past_key_values=cache, use_cache=True)
+        return out.logits[0, -1].float().cpu().numpy(), out.past_key_values
+
+    texts = [tokenizer.decode([i]) for i in range(len(tokenizer))]
+    stops = {tokenizer.eos_token_id} if tokenizer.eos_token_id is not None else set()
+    encode = lambda text: tokenizer.encode(text, add_special_tokens=False) or [tokenizer.eos_token_id]
+    return TokenLM(name, next_logits, encode, texts, stops)
+
+
+def onnx_token_lm(name: str, folder: str, fim: bool) -> TokenLM:
+    """A tinynext v2 export: graph.onnx + weights.f32 (compressed, expanded) or model.onnx, plus tokenizer.json."""
+    import onnxruntime as ort
+    from tokenizers import Tokenizer
+
+    folder = Path(folder)
+    graph = folder / "graph.onnx" if (folder / "graph.onnx").exists() else folder / "model.onnx"
+    session = ort.InferenceSession(str(graph), providers=["CPUExecutionProvider"])
+    config = json.loads((folder / "config.json").read_text())
+    tokenizer = Tokenizer.from_file(str(folder / "tokenizer.json"))
+    empty = np.zeros((config["layers"], 1, config["heads"], 0, config["dim"] // config["heads"]), np.float32)
+
+    def next_logits(ids, cache):
+        past_k, past_v = cache if cache else (empty, empty)
+        logits, k, v = session.run(None, {"input_ids": np.array([ids], np.int64), "past_k": past_k, "past_v": past_v})
+        return logits[0], (k, v)
+
+    special = len(config["special_tokens"])
+    texts = [tokenizer.decode([i], skip_special_tokens=False) if i >= special else "" for i in range(config["vocab"])]
+    encode = lambda text: tokenizer.encode(text).ids or [0]
+    prompt = (lambda prefix, suffix: prefix if not suffix else [2, *suffix, 1, *prefix, 3]) if fim else None
+    return TokenLM(name, next_logits, encode, texts, set(range(special)), prompt=prompt)
+
+
+# ---- Scoring ---------------------------------------------------------------------------------------------------
+
+
+def word_by_word(suggestion: str, truth: str) -> int:
+    """Characters you'd keep accepting a word at a time until a word is wrong."""
+    kept = 0
+    for match in re.finditer(r"\s*\S+", suggestion):
+        if truth.startswith(suggestion[: match.end()]):
+            kept = match.end()
+        else:
+            break
+    return kept
+
+
+def score(model: Model, all_pauses: list[Pause]) -> dict:
+    results: dict = {"model": model.name, "sees_after": model.sees_after, "cutoffs": {}, "examples": [], "latency_ms": {}}
+    latencies: list[float] = []
+    raw = []
+    for index, pause in enumerate(all_pauses):
+        trajectory = model.suggest(pause.before, pause.after if model.sees_after else "")
+        latencies.append(trajectory.milliseconds)
+        raw.append((pause, trajectory))
+        if index % 500 == 0:
+            print(f"  {model.name}: {index}/{len(all_pauses)} pauses, median {statistics.median(latencies):.1f} ms", flush=True)
+
+    for cutoff in CUTOFFS:
+        table: dict = {}
+        for domain in DOMAINS:
+            for scenario in SCENARIOS:
+                rows = [(p, t.at(cutoff)) for p, t in raw if p.domain == domain and p.scenario == scenario]
+                if not rows:
+                    continue
+                shown = [(p, s) for p, s in rows if len(s) >= MIN_SHOWN_CHARS]
+                right = [(p, s) for p, s in shown if p.truth.startswith(s)]
+                multi = [(p, s) for p, s in shown if len(s.split()) >= 2]
+                multi_right = [(p, s) for p, s in multi if p.truth.startswith(s)]
+                typed = len(rows) * PAUSE_EVERY
+                table[f"{domain}/{scenario}"] = {
+                    "pauses": len(rows),
+                    "shown": round(len(shown) / len(rows), 4),
+                    "right": round(len(right) / max(1, len(shown)), 4),
+                    "avg_right_chars": round(sum(len(s) for _, s in right) / max(1, len(right)), 2),
+                    "multiword_shown": round(len(multi) / len(rows), 4),
+                    "multiword_right": round(len(multi_right) / max(1, len(multi)), 4),
+                    "saved_with_tab": round(sum(len(s) for _, s in right) / typed, 4),
+                    "saved_word_by_word": round(sum(word_by_word(s, p.truth) for p, s in shown) / typed, 4),
+                }
+        results["cutoffs"][str(cutoff)] = table
+
+    for pause, trajectory in raw[:: max(1, len(raw) // 60)]:
+        suggestion = trajectory.at(MAIN_CUTOFF)
+        if len(suggestion) >= MIN_SHOWN_CHARS:
+            results["examples"].append({
+                "domain": pause.domain, "scenario": pause.scenario, "before": pause.before[-60:],
+                "suggestion": suggestion, "truth": pause.truth[: len(suggestion) + 12], "right": pause.truth.startswith(suggestion),
+            })
+    results["latency_ms"] = {
+        "p50": round(statistics.median(latencies), 2),
+        "p95": round(float(np.percentile(latencies, 95)), 2),
+        "note": "Measured in this eval process, not in a browser.",
+    }
+    return results
+
+
+# ---- Report ----------------------------------------------------------------------------------------------------
+
+
+def table(paths: list[Path]) -> str:
+    runs = [json.loads(p.read_text()) for p in paths]
+    lines = []
+    for scenario in SCENARIOS:
+        lines.append(f"\n### {scenario.capitalize()}, {int(MAIN_CUTOFF * 100)}% cutoff, suggestions of {MIN_SHOWN_CHARS}+ characters\n")
+        lines.append("| Model | Domain | Shown | Right | Avg right | 2+ words shown | 2+ words right | Saved (Tab) | Saved (word by word) |")
+        lines.append("|---|---|---|---|---|---|---|---|---|")
+        for run in runs:
+            for domain in DOMAINS:
+                m = run["cutoffs"][str(MAIN_CUTOFF)].get(f"{domain}/{scenario}")
+                if not m:
+                    continue
+                lines.append(
+                    f"| {run['model']} | {domain} | {m['shown']:.0%} | {m['right']:.0%} | {m['avg_right_chars']:.1f} ch "
+                    f"| {m['multiword_shown']:.0%} | {m['multiword_right']:.0%} | {m['saved_with_tab']:.1%} | {m['saved_word_by_word']:.1%} |"
+                )
+    lines.append("\n| Model | Latency p50 | p95 |\n|---|---|---|")
+    for run in runs:
+        lines.append(f"| {run['model']} | {run['latency_ms']['p50']} ms | {run['latency_ms']['p95']} ms |")
+    return "\n".join(lines)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", help="Name for the metrics file")
+    parser.add_argument("--onnx", help="A character-level GRU exported like enron-autocomplete")
+    parser.add_argument("--hf", help="A Hugging Face causal LM to evaluate, e.g. HuggingFaceTB/SmolLM2-135M")
+    parser.add_argument("--onnx-lm", help="Folder with a tinynext v2 ONNX export, config.json and tokenizer.json")
+    parser.add_argument("--no-fim", action="store_true", help="With --onnx-lm: the model wasn't trained to fill in the middle")
+    parser.add_argument("--max-pauses", type=int, default=1000, help="Per domain and scenario")
+    parser.add_argument("--table", action="store_true", help="Print a comparison of every metrics/*.json")
+    args = parser.parse_args()
+
+    if args.table:
+        print(table(sorted(METRICS.glob("*.json"))))
+        return
+    if args.onnx:
+        model: Model = CharGRU(args.model, args.onnx)
+    elif args.hf:
+        model = hf_causal_lm(args.model, args.hf)
+    elif args.onnx_lm:
+        model = onnx_token_lm(args.model, args.onnx_lm, fim=not args.no_fim)
+    else:
+        parser.error("Pass --onnx, --hf or --table. train.py evaluates its own model through score().")
+    all_pauses = pauses(load_eval_set(), args.max_pauses)
+    print(f"{len(all_pauses)} pauses", flush=True)
+    results = score(model, all_pauses)
+    METRICS.mkdir(exist_ok=True)
+    (METRICS / f"{args.model}.json").write_text(json.dumps(results, indent=1))
+    print(table([METRICS / f"{args.model}.json"]))
+
+
+if __name__ == "__main__":
+    main()

v2/experiments/tinynext-v2-3b/metrics/tinynext-v2-3b.json

@@ -0,0 +1,371 @@
+{
+ "model": "tinynext-v2-3b",
+ "sees_after": true,
+ "cutoffs": {
+  "0.3": {
+   "articles/typing": {
+    "pauses": 1000,
+    "shown": 0.204,
+    "right": 0.5882,
+    "avg_right_chars": 6.04,
+    "multiword_shown": 0.086,
+    "multiword_right": 0.5465,
+    "saved_with_tab": 0.0604,
+    "saved_word_by_word": 0.0679
+   },
+   "articles/editing": {
+    "pauses": 1000,
+    "shown": 0.152,
+    "right": 0.6382,
+    "avg_right_chars": 5.86,
+    "multiword_shown": 0.062,
+    "multiword_right": 0.629,
+    "saved_with_tab": 0.0473,
+    "saved_word_by_word": 0.0516
+   },
+   "messages/typing": {
+    "pauses": 1000,
+    "shown": 0.324,
+    "right": 0.5463,
+    "avg_right_chars": 7.8,
+    "multiword_shown": 0.18,
+    "multiword_right": 0.5333,
+    "saved_with_tab": 0.115,
+    "saved_word_by_word": 0.1404
+   },
+   "messages/editing": {
+    "pauses": 1000,
+    "shown": 0.317,
+    "right": 0.5584,
+    "avg_right_chars": 7.47,
+    "multiword_shown": 0.173,
+    "multiword_right": 0.5376,
+    "saved_with_tab": 0.1102,
+    "saved_word_by_word": 0.1328
+   },
+   "email/typing": {
+    "pauses": 1000,
+    "shown": 0.234,
+    "right": 0.5598,
+    "avg_right_chars": 6.03,
+    "multiword_shown": 0.087,
+    "multiword_right": 0.4023,
+    "saved_with_tab": 0.0658,
+    "saved_word_by_word": 0.0746
+   },
+   "email/editing": {
+    "pauses": 1000,
+    "shown": 0.213,
+    "right": 0.5915,
+    "avg_right_chars": 5.63,
+    "multiword_shown": 0.078,
+    "multiword_right": 0.4231,
+    "saved_with_tab": 0.0591,
+    "saved_word_by_word": 0.0653
+   }
+  },
+  "0.4": {
+   "articles/typing": {
+    "pauses": 1000,
+    "shown": 0.149,
+    "right": 0.698,
+    "avg_right_chars": 6.12,
+    "multiword_shown": 0.057,
+    "multiword_right": 0.614,
+    "saved_with_tab": 0.0531,
+    "saved_word_by_word": 0.0563
+   },
+   "articles/editing": {
+    "pauses": 1000,
+    "shown": 0.107,
+    "right": 0.7103,
+    "avg_right_chars": 5.72,
+    "multiword_shown": 0.034,
+    "multiword_right": 0.6765,
+    "saved_with_tab": 0.0362,
+    "saved_word_by_word": 0.0384
+   },
+   "messages/typing": {
+    "pauses": 1000,
+    "shown": 0.251,
+    "right": 0.6255,
+    "avg_right_chars": 7.7,
+    "multiword_shown": 0.133,
+    "multiword_right": 0.6015,
+    "saved_with_tab": 0.1008,
+    "saved_word_by_word": 0.1159
+   },
+   "messages/editing": {
+    "pauses": 1000,
+    "shown": 0.239,
+    "right": 0.6569,
+    "avg_right_chars": 7.29,
+    "multiword_shown": 0.127,
+    "multiword_right": 0.6378,
+    "saved_with_tab": 0.0954,
+    "saved_word_by_word": 0.1078
+   },
+   "email/typing": {
+    "pauses": 1000,
+    "shown": 0.162,
+    "right": 0.679,
+    "avg_right_chars": 5.84,
+    "multiword_shown": 0.049,
+    "multiword_right": 0.4898,
+    "saved_with_tab": 0.0535,
+    "saved_word_by_word": 0.0575
+   },
+   "email/editing": {
+    "pauses": 1000,
+    "shown": 0.165,
+    "right": 0.6727,
+    "avg_right_chars": 5.56,
+    "multiword_shown": 0.055,
+    "multiword_right": 0.5818,
+    "saved_with_tab": 0.0514,
+    "saved_word_by_word": 0.0545
+   }
+  },
+  "0.5": {
+   "articles/typing": {
+    "pauses": 1000,
+    "shown": 0.101,
+    "right": 0.7822,
+    "avg_right_chars": 5.86,
+    "multiword_shown": 0.035,
+    "multiword_right": 0.7143,
+    "saved_with_tab": 0.0386,
+    "saved_word_by_word": 0.04
+   },
+   "articles/editing": {
+    "pauses": 1000,
+    "shown": 0.08,
+    "right": 0.825,
+    "avg_right_chars": 5.61,
+    "multiword_shown": 0.024,
+    "multiword_right": 0.8333,
+    "saved_with_tab": 0.0308,
+    "saved_word_by_word": 0.0317
+   },
+   "messages/typing": {
+    "pauses": 1000,
+    "shown": 0.193,
+    "right": 0.7461,
+    "avg_right_chars": 7.31,
+    "multiword_shown": 0.094,
+    "multiword_right": 0.7021,
+    "saved_with_tab": 0.0877,
+    "saved_word_by_word": 0.0969
+   },
+   "messages/editing": {
+    "pauses": 1000,
+    "shown": 0.175,
+    "right": 0.76,
+    "avg_right_chars": 6.87,
+    "multiword_shown": 0.091,
+    "multiword_right": 0.7033,
+    "saved_with_tab": 0.0762,
+    "saved_word_by_word": 0.0843
+   },
+   "email/typing": {
+    "pauses": 1000,
+    "shown": 0.112,
+    "right": 0.7768,
+    "avg_right_chars": 5.54,
+    "multiword_shown": 0.035,
+    "multiword_right": 0.5143,
+    "saved_with_tab": 0.0402,
+    "saved_word_by_word": 0.0427
+   },
+   "email/editing": {
+    "pauses": 1000,
+    "shown": 0.118,
+    "right": 0.7797,
+    "avg_right_chars": 5.5,
+    "multiword_shown": 0.031,
+    "multiword_right": 0.7097,
+    "saved_with_tab": 0.0422,
+    "saved_word_by_word": 0.0435
+   }
+  },
+  "0.6": {
+   "articles/typing": {
+    "pauses": 1000,
+    "shown": 0.076,
+    "right": 0.8026,
+    "avg_right_chars": 5.77,
+    "multiword_shown": 0.026,
+    "multiword_right": 0.7692,
+    "saved_with_tab": 0.0293,
+    "saved_word_by_word": 0.0303
+   },
+   "articles/editing": {
+    "pauses": 1000,
+    "shown": 0.059,
+    "right": 0.8983,
+    "avg_right_chars": 5.45,
+    "multiword_shown": 0.014,
+    "multiword_right": 0.7857,
+    "saved_with_tab": 0.0241,
+    "saved_word_by_word": 0.0244
+   },
+   "messages/typing": {
+    "pauses": 1000,
+    "shown": 0.152,
+    "right": 0.8421,
+    "avg_right_chars": 7.22,
+    "multiword_shown": 0.068,
+    "multiword_right": 0.8382,
+    "saved_with_tab": 0.077,
+    "saved_word_by_word": 0.0795
+   },
+   "messages/editing": {
+    "pauses": 1000,
+    "shown": 0.141,
+    "right": 0.8369,
+    "avg_right_chars": 6.39,
+    "multiword_shown": 0.065,
+    "multiword_right": 0.7538,
+    "saved_with_tab": 0.0628,
+    "saved_word_by_word": 0.0669
+   },
+   "email/typing": {
+    "pauses": 1000,
+    "shown": 0.086,
+    "right": 0.8256,
+    "avg_right_chars": 5.23,
+    "multiword_shown": 0.015,
+    "multiword_right": 0.5333,
+    "saved_with_tab": 0.0309,
+    "saved_word_by_word": 0.0322
+   },
+   "email/editing": {
+    "pauses": 1000,
+    "shown": 0.084,
+    "right": 0.8452,
+    "avg_right_chars": 5.34,
+    "multiword_shown": 0.016,
+    "multiword_right": 0.75,
+    "saved_with_tab": 0.0316,
+    "saved_word_by_word": 0.0319
+   }
+  },
+  "0.7": {
+   "articles/typing": {
+    "pauses": 1000,
+    "shown": 0.052,
+    "right": 0.8654,
+    "avg_right_chars": 6.09,
+    "multiword_shown": 0.017,
+    "multiword_right": 0.8235,
+    "saved_with_tab": 0.0228,
+    "saved_word_by_word": 0.0232
+   },
+   "articles/editing": {
+    "pauses": 1000,
+    "shown": 0.043,
+    "right": 0.9302,
+    "avg_right_chars": 5.25,
+    "multiword_shown": 0.01,
+    "multiword_right": 0.9,
+    "saved_with_tab": 0.0175,
+    "saved_word_by_word": 0.0177
+   },
+   "messages/typing": {
+    "pauses": 1000,
+    "shown": 0.117,
+    "right": 0.8718,
+    "avg_right_chars": 6.64,
+    "multiword_shown": 0.045,
+    "multiword_right": 0.8889,
+    "saved_with_tab": 0.0564,
+    "saved_word_by_word": 0.0579
+   },
+   "messages/editing": {
+    "pauses": 1000,
+    "shown": 0.105,
+    "right": 0.8857,
+    "avg_right_chars": 5.98,
+    "multiword_shown": 0.038,
+    "multiword_right": 0.8684,
+    "saved_with_tab": 0.0463,
+    "saved_word_by_word": 0.0476
+   },
+   "email/typing": {
+    "pauses": 1000,
+    "shown": 0.058,
+    "right": 0.8103,
+    "avg_right_chars": 5.19,
+    "multiword_shown": 0.007,
+    "multiword_right": 0.4286,
+    "saved_with_tab": 0.0203,
+    "saved_word_by_word": 0.0206
+   },
+   "email/editing": {
+    "pauses": 1000,
+    "shown": 0.066,
+    "right": 0.8788,
+    "avg_right_chars": 4.98,
+    "multiword_shown": 0.008,
+    "multiword_right": 0.5,
+    "saved_with_tab": 0.0241,
+    "saved_word_by_word": 0.0244
+   }
+  }
+ },
+ "examples": [
+  {
+   "domain": "articles",
+   "scenario": "editing",
+   "before": " the Spaghetti Western film genre) and Hate Boat, were publi",
+   "suggestion": "shed",
+   "truth": "shed by Comico C",
+   "right": true
+  },
+  {
+   "domain": "articles",
+   "scenario": "editing",
+   "before": "rked on a tour of Canada with I Hate Sally, Hiretsukan, GFK,",
+   "suggestion": " and",
+   "truth": " and DeadPoets a",
+   "right": true
+  },
+  {
+   "domain": "messages",
+   "scenario": "typing",
+   "before": " be completely benevolent and will always act in our best in",
+   "suggestion": "terests.",
+   "truth": "terests, ensuring th",
+   "right": false
+  },
+  {
+   "domain": "messages",
+   "scenario": "typing",
+   "before": "onsidered new. Those decisions make the graphs as neat as po",
+   "suggestion": "ssible.",
+   "truth": "ssible, and it is t",
+   "right": false
+  },
+  {
+   "domain": "messages",
+   "scenario": "editing",
+   "before": "e optimal level of production and pricing for businesses. L'",
+   "suggestion": "Hpital's Rule can",
+   "truth": "Hpital's Rule can help analyz",
+   "right": true
+  },
+  {
+   "domain": "email",
+   "scenario": "editing",
+   "before": "ured products and marketing of electricity for the United St",
+   "suggestion": "ates",
+   "truth": "ates ? Actively ",
+   "right": true
+  }
+ ],
+ "latency_ms": {
+  "p50": 15.38,
+  "p95": 42.67,
+  "note": "Measured in this eval process, not in a browser."
+ }
+}

v2/experiments/tinynext-v2-3b/model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d720a8d92844daca8da7820debd0be2c22d064b33de16ccbcfa69adf79307f13
+size 60813561

v2/experiments/tinynext-v2-3b/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0ffd7eaaee87262d76e9b401514eee2449e72b7aae899579b227c114d4a37366
+size 50161328

v2/experiments/tinynext-v2-3b/model_int8.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:34f405276119c07d523af9a5e43354ab50acbaa0e2be167166320401a3d35a2c
+size 15638872

v2/experiments/tinynext-v2-3b/report.json

@@ -0,0 +1,2390 @@
+{
+ "args": {
+  "vocab": 8192,
+  "dim": 320,
+  "layers": 8,
+  "heads": 5,
+  "ffn": 864,
+  "seq": 512,
+  "tokens": 3000000000.0,
+  "val_tokens": 4000000.0,
+  "epochs": 1.0,
+  "batch": 512,
+  "lr": 0.003,
+  "fim_rate": 0.5,
+  "mix": "fineweb=0.6,smoltalk=0.3,enron=0.1",
+  "fineweb_files": 4,
+  "smoltalk_files": 9,
+  "tokenizer_docs": 200000,
+  "eval_pauses": 1000,
+  "log_every": 50,
+  "repo": "cfahlgren1/tinynext",
+  "folder": "v2/experiments/tinynext-v2-3b",
+  "eval_folder": "v2",
+  "name": "tinynext-v2-3b",
+  "trackio_space": "cfahlgren1/tinynext-trackio",
+  "stage": "train",
+  "data_repo": "cfahlgren1/tinynext-data",
+  "smoke": false,
+  "no_push": false
+ },
+ "parameters": 12539200,
+ "data": {
+  "fineweb": {
+   "docs": 2122073,
+   "train_tokens": 1800000000
+  },
+  "smoltalk": {
+   "docs": 1041801,
+   "train_tokens": 900000000
+  },
+  "enron": {
+   "docs": 268120,
+   "train_tokens": 300000000
+  }
+ },
+ "data_minutes": 0.2,
+ "train_minutes": 55.8,
+ "history": [
+  {
+   "step": 0,
+   "train_loss": 9.0702,
+   "lr": 1.3e-05,
+   "tokens_per_second": 12800,
+   "minutes": 0.34,
+   "mfu_percent": 0.3,
+   "gpu_util_percent": 98,
+   "gpu_memory_gb": 76.5,
+   "val_loss": 9.0663
+  },
+  {
+   "step": 50,
+   "train_loss": 7.215,
+   "lr": 0.000671,
+   "tokens_per_second": 930958,
+   "minutes": 0.62,
+   "mfu_percent": 22.4,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 100,
+   "train_loss": 6.1306,
+   "lr": 0.001329,
+   "tokens_per_second": 925320,
+   "minutes": 0.86,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 150,
+   "train_loss": 5.5618,
+   "lr": 0.001986,
+   "tokens_per_second": 924266,
+   "minutes": 1.1,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 200,
+   "train_loss": 5.1433,
+   "lr": 0.002643,
+   "tokens_per_second": 923447,
+   "minutes": 1.33,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 250,
+   "train_loss": 4.6259,
+   "lr": 0.002997,
+   "tokens_per_second": 923503,
+   "minutes": 1.57,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 300,
+   "train_loss": 4.2702,
+   "lr": 0.002995,
+   "tokens_per_second": 923163,
+   "minutes": 1.81,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 350,
+   "train_loss": 4.0271,
+   "lr": 0.002994,
+   "tokens_per_second": 922754,
+   "minutes": 2.04,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 400,
+   "train_loss": 3.8697,
+   "lr": 0.002992,
+   "tokens_per_second": 922255,
+   "minutes": 2.28,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 450,
+   "train_loss": 3.6918,
+   "lr": 0.00299,
+   "tokens_per_second": 921900,
+   "minutes": 2.52,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 500,
+   "train_loss": 3.6502,
+   "lr": 0.002987,
+   "tokens_per_second": 921603,
+   "minutes": 2.75,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1,
+   "val_loss": 3.6813
+  },
+  {
+   "step": 550,
+   "train_loss": 3.6069,
+   "lr": 0.002985,
+   "tokens_per_second": 922139,
+   "minutes": 3.04,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 600,
+   "train_loss": 3.6009,
+   "lr": 0.002982,
+   "tokens_per_second": 921841,
+   "minutes": 3.28,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 650,
+   "train_loss": 3.5049,
+   "lr": 0.002978,
+   "tokens_per_second": 921515,
+   "minutes": 3.51,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 700,
+   "train_loss": 3.4728,
+   "lr": 0.002975,
+   "tokens_per_second": 921918,
+   "minutes": 3.75,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 750,
+   "train_loss": 3.4877,
+   "lr": 0.002971,
+   "tokens_per_second": 921386,
+   "minutes": 3.99,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 800,
+   "train_loss": 3.4264,
+   "lr": 0.002967,
+   "tokens_per_second": 921002,
+   "minutes": 4.22,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 850,
+   "train_loss": 3.3612,
+   "lr": 0.002963,
+   "tokens_per_second": 920889,
+   "minutes": 4.46,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 900,
+   "train_loss": 3.3516,
+   "lr": 0.002959,
+   "tokens_per_second": 920832,
+   "minutes": 4.7,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 950,
+   "train_loss": 3.3742,
+   "lr": 0.002954,
+   "tokens_per_second": 920862,
+   "minutes": 4.94,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 1000,
+   "train_loss": 3.3632,
+   "lr": 0.002949,
+   "tokens_per_second": 920554,
+   "minutes": 5.17,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1,
+   "val_loss": 3.3796
+  },
+  {
+   "step": 1050,
+   "train_loss": 3.3061,
+   "lr": 0.002944,
+   "tokens_per_second": 920467,
+   "minutes": 5.46,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 1100,
+   "train_loss": 3.3351,
+   "lr": 0.002939,
+   "tokens_per_second": 921248,
+   "minutes": 5.69,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 1150,
+   "train_loss": 3.3238,
+   "lr": 0.002933,
+   "tokens_per_second": 921674,
+   "minutes": 5.93,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 1200,
+   "train_loss": 3.3208,
+   "lr": 0.002927,
+   "tokens_per_second": 921693,
+   "minutes": 6.17,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 1250,
+   "train_loss": 3.286,
+   "lr": 0.002921,
+   "tokens_per_second": 921893,
+   "minutes": 6.41,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 1300,
+   "train_loss": 3.2941,
+   "lr": 0.002914,
+   "tokens_per_second": 922496,
+   "minutes": 6.64,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 1350,
+   "train_loss": 3.2475,
+   "lr": 0.002908,
+   "tokens_per_second": 922535,
+   "minutes": 6.88,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 1400,
+   "train_loss": 3.21,
+   "lr": 0.002901,
+   "tokens_per_second": 921917,
+   "minutes": 7.12,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 1450,
+   "train_loss": 3.2507,
+   "lr": 0.002894,
+   "tokens_per_second": 921722,
+   "minutes": 7.35,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 1500,
+   "train_loss": 3.2548,
+   "lr": 0.002887,
+   "tokens_per_second": 920926,
+   "minutes": 7.59,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1,
+   "val_loss": 3.2689
+  },
+  {
+   "step": 1550,
+   "train_loss": 3.2187,
+   "lr": 0.002879,
+   "tokens_per_second": 920076,
+   "minutes": 7.87,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 1600,
+   "train_loss": 3.2435,
+   "lr": 0.002871,
+   "tokens_per_second": 920390,
+   "minutes": 8.11,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 1650,
+   "train_loss": 3.1958,
+   "lr": 0.002863,
+   "tokens_per_second": 920901,
+   "minutes": 8.35,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 1700,
+   "train_loss": 3.1797,
+   "lr": 0.002855,
+   "tokens_per_second": 920930,
+   "minutes": 8.59,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 1750,
+   "train_loss": 3.2016,
+   "lr": 0.002846,
+   "tokens_per_second": 921134,
+   "minutes": 8.82,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 1800,
+   "train_loss": 3.1736,
+   "lr": 0.002838,
+   "tokens_per_second": 922208,
+   "minutes": 9.06,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 1850,
+   "train_loss": 3.1912,
+   "lr": 0.002829,
+   "tokens_per_second": 922453,
+   "minutes": 9.3,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 1900,
+   "train_loss": 3.1827,
+   "lr": 0.00282,
+   "tokens_per_second": 922660,
+   "minutes": 9.53,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 1950,
+   "train_loss": 3.202,
+   "lr": 0.00281,
+   "tokens_per_second": 922623,
+   "minutes": 9.77,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 2000,
+   "train_loss": 3.2288,
+   "lr": 0.002801,
+   "tokens_per_second": 922474,
+   "minutes": 10.01,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1,
+   "val_loss": 3.208
+  },
+  {
+   "step": 2050,
+   "train_loss": 3.1591,
+   "lr": 0.002791,
+   "tokens_per_second": 922197,
+   "minutes": 10.29,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 2100,
+   "train_loss": 3.2074,
+   "lr": 0.002781,
+   "tokens_per_second": 922026,
+   "minutes": 10.53,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 2150,
+   "train_loss": 3.1784,
+   "lr": 0.002771,
+   "tokens_per_second": 922145,
+   "minutes": 10.77,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 2200,
+   "train_loss": 3.1416,
+   "lr": 0.00276,
+   "tokens_per_second": 921690,
+   "minutes": 11.0,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 2250,
+   "train_loss": 3.1531,
+   "lr": 0.002749,
+   "tokens_per_second": 921458,
+   "minutes": 11.24,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 2300,
+   "train_loss": 3.1148,
+   "lr": 0.002738,
+   "tokens_per_second": 920727,
+   "minutes": 11.48,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 2350,
+   "train_loss": 3.1005,
+   "lr": 0.002727,
+   "tokens_per_second": 921746,
+   "minutes": 11.71,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 2400,
+   "train_loss": 3.2112,
+   "lr": 0.002716,
+   "tokens_per_second": 922555,
+   "minutes": 11.95,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 2450,
+   "train_loss": 3.1437,
+   "lr": 0.002705,
+   "tokens_per_second": 922219,
+   "minutes": 12.19,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 2500,
+   "train_loss": 3.114,
+   "lr": 0.002693,
+   "tokens_per_second": 922404,
+   "minutes": 12.42,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1,
+   "val_loss": 3.1664
+  },
+  {
+   "step": 2550,
+   "train_loss": 3.1538,
+   "lr": 0.002681,
+   "tokens_per_second": 922093,
+   "minutes": 12.71,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 2600,
+   "train_loss": 3.1874,
+   "lr": 0.002669,
+   "tokens_per_second": 921696,
+   "minutes": 12.95,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 2650,
+   "train_loss": 3.1451,
+   "lr": 0.002657,
+   "tokens_per_second": 920726,
+   "minutes": 13.18,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 2700,
+   "train_loss": 3.1138,
+   "lr": 0.002644,
+   "tokens_per_second": 920915,
+   "minutes": 13.42,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 2750,
+   "train_loss": 3.1194,
+   "lr": 0.002632,
+   "tokens_per_second": 920910,
+   "minutes": 13.66,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 2800,
+   "train_loss": 3.0835,
+   "lr": 0.002619,
+   "tokens_per_second": 920819,
+   "minutes": 13.89,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 2850,
+   "train_loss": 3.1095,
+   "lr": 0.002606,
+   "tokens_per_second": 920420,
+   "minutes": 14.13,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 2900,
+   "train_loss": 3.1067,
+   "lr": 0.002593,
+   "tokens_per_second": 920261,
+   "minutes": 14.37,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 2950,
+   "train_loss": 3.0488,
+   "lr": 0.002579,
+   "tokens_per_second": 921100,
+   "minutes": 14.61,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 3000,
+   "train_loss": 3.1089,
+   "lr": 0.002566,
+   "tokens_per_second": 921915,
+   "minutes": 14.84,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1,
+   "val_loss": 3.1334
+  },
+  {
+   "step": 3050,
+   "train_loss": 3.0736,
+   "lr": 0.002552,
+   "tokens_per_second": 922085,
+   "minutes": 15.13,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 3100,
+   "train_loss": 3.0915,
+   "lr": 0.002538,
+   "tokens_per_second": 921538,
+   "minutes": 15.36,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 3150,
+   "train_loss": 3.0751,
+   "lr": 0.002524,
+   "tokens_per_second": 921498,
+   "minutes": 15.6,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 3200,
+   "train_loss": 3.1093,
+   "lr": 0.00251,
+   "tokens_per_second": 920747,
+   "minutes": 15.84,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 3250,
+   "train_loss": 3.1425,
+   "lr": 0.002495,
+   "tokens_per_second": 920376,
+   "minutes": 16.08,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 3300,
+   "train_loss": 3.0762,
+   "lr": 0.002481,
+   "tokens_per_second": 920083,
+   "minutes": 16.31,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 3350,
+   "train_loss": 3.1273,
+   "lr": 0.002466,
+   "tokens_per_second": 919889,
+   "minutes": 16.55,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 3400,
+   "train_loss": 3.1,
+   "lr": 0.002451,
+   "tokens_per_second": 921166,
+   "minutes": 16.79,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 3450,
+   "train_loss": 3.1428,
+   "lr": 0.002436,
+   "tokens_per_second": 921093,
+   "minutes": 17.02,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 3500,
+   "train_loss": 3.0703,
+   "lr": 0.002421,
+   "tokens_per_second": 920938,
+   "minutes": 17.26,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1,
+   "val_loss": 3.1114
+  },
+  {
+   "step": 3550,
+   "train_loss": 3.1048,
+   "lr": 0.002406,
+   "tokens_per_second": 921446,
+   "minutes": 17.55,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 3600,
+   "train_loss": 3.0576,
+   "lr": 0.00239,
+   "tokens_per_second": 921917,
+   "minutes": 17.78,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 3650,
+   "train_loss": 3.0815,
+   "lr": 0.002375,
+   "tokens_per_second": 921369,
+   "minutes": 18.02,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 3700,
+   "train_loss": 3.102,
+   "lr": 0.002359,
+   "tokens_per_second": 921425,
+   "minutes": 18.26,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 3750,
+   "train_loss": 3.0362,
+   "lr": 0.002343,
+   "tokens_per_second": 921350,
+   "minutes": 18.49,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 3800,
+   "train_loss": 3.0936,
+   "lr": 0.002327,
+   "tokens_per_second": 920994,
+   "minutes": 18.73,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 3850,
+   "train_loss": 3.0723,
+   "lr": 0.002311,
+   "tokens_per_second": 920745,
+   "minutes": 18.97,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 3900,
+   "train_loss": 3.0621,
+   "lr": 0.002295,
+   "tokens_per_second": 920434,
+   "minutes": 19.21,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 3950,
+   "train_loss": 3.1088,
+   "lr": 0.002278,
+   "tokens_per_second": 919905,
+   "minutes": 19.44,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 4000,
+   "train_loss": 3.019,
+   "lr": 0.002262,
+   "tokens_per_second": 919709,
+   "minutes": 19.68,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1,
+   "val_loss": 3.0871
+  },
+  {
+   "step": 4050,
+   "train_loss": 3.0537,
+   "lr": 0.002245,
+   "tokens_per_second": 919817,
+   "minutes": 19.97,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 4100,
+   "train_loss": 3.0237,
+   "lr": 0.002228,
+   "tokens_per_second": 919826,
+   "minutes": 20.2,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 4150,
+   "train_loss": 3.0252,
+   "lr": 0.002212,
+   "tokens_per_second": 920702,
+   "minutes": 20.44,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 4200,
+   "train_loss": 2.9956,
+   "lr": 0.002195,
+   "tokens_per_second": 920999,
+   "minutes": 20.68,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 4250,
+   "train_loss": 3.0576,
+   "lr": 0.002178,
+   "tokens_per_second": 920918,
+   "minutes": 20.92,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 4300,
+   "train_loss": 3.0944,
+   "lr": 0.00216,
+   "tokens_per_second": 921070,
+   "minutes": 21.15,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 4350,
+   "train_loss": 3.0655,
+   "lr": 0.002143,
+   "tokens_per_second": 921341,
+   "minutes": 21.39,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 4400,
+   "train_loss": 3.0756,
+   "lr": 0.002126,
+   "tokens_per_second": 920880,
+   "minutes": 21.63,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 4450,
+   "train_loss": 3.0513,
+   "lr": 0.002108,
+   "tokens_per_second": 921004,
+   "minutes": 21.86,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 4500,
+   "train_loss": 3.0553,
+   "lr": 0.002091,
+   "tokens_per_second": 921256,
+   "minutes": 22.1,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1,
+   "val_loss": 3.0698
+  },
+  {
+   "step": 4550,
+   "train_loss": 2.967,
+   "lr": 0.002073,
+   "tokens_per_second": 920961,
+   "minutes": 22.39,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 4600,
+   "train_loss": 3.0567,
+   "lr": 0.002056,
+   "tokens_per_second": 920043,
+   "minutes": 22.62,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 4650,
+   "train_loss": 3.0559,
+   "lr": 0.002038,
+   "tokens_per_second": 919645,
+   "minutes": 22.86,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 4700,
+   "train_loss": 3.0584,
+   "lr": 0.00202,
+   "tokens_per_second": 920036,
+   "minutes": 23.1,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 4750,
+   "train_loss": 3.0321,
+   "lr": 0.002002,
+   "tokens_per_second": 920584,
+   "minutes": 23.34,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 4800,
+   "train_loss": 2.9678,
+   "lr": 0.001984,
+   "tokens_per_second": 920212,
+   "minutes": 23.57,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 4850,
+   "train_loss": 2.9969,
+   "lr": 0.001966,
+   "tokens_per_second": 921129,
+   "minutes": 23.81,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 4900,
+   "train_loss": 3.0317,
+   "lr": 0.001948,
+   "tokens_per_second": 921336,
+   "minutes": 24.05,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 4950,
+   "train_loss": 3.0314,
+   "lr": 0.00193,
+   "tokens_per_second": 921674,
+   "minutes": 24.28,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 5000,
+   "train_loss": 3.0467,
+   "lr": 0.001912,
+   "tokens_per_second": 922491,
+   "minutes": 24.52,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1,
+   "val_loss": 3.0507
+  },
+  {
+   "step": 5050,
+   "train_loss": 3.0208,
+   "lr": 0.001894,
+   "tokens_per_second": 922889,
+   "minutes": 24.8,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 5100,
+   "train_loss": 2.9957,
+   "lr": 0.001875,
+   "tokens_per_second": 922887,
+   "minutes": 25.04,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 5150,
+   "train_loss": 3.0355,
+   "lr": 0.001857,
+   "tokens_per_second": 922331,
+   "minutes": 25.28,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 5200,
+   "train_loss": 3.0702,
+   "lr": 0.001839,
+   "tokens_per_second": 922142,
+   "minutes": 25.51,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 5250,
+   "train_loss": 3.0055,
+   "lr": 0.00182,
+   "tokens_per_second": 922504,
+   "minutes": 25.75,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 5300,
+   "train_loss": 2.9704,
+   "lr": 0.001802,
+   "tokens_per_second": 922274,
+   "minutes": 25.99,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 5350,
+   "train_loss": 2.9492,
+   "lr": 0.001783,
+   "tokens_per_second": 921682,
+   "minutes": 26.23,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 5400,
+   "train_loss": 3.0025,
+   "lr": 0.001765,
+   "tokens_per_second": 921085,
+   "minutes": 26.46,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 5450,
+   "train_loss": 2.9997,
+   "lr": 0.001746,
+   "tokens_per_second": 921161,
+   "minutes": 26.7,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 5500,
+   "train_loss": 3.012,
+   "lr": 0.001728,
+   "tokens_per_second": 921331,
+   "minutes": 26.94,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1,
+   "val_loss": 3.0312
+  },
+  {
+   "step": 5550,
+   "train_loss": 3.0025,
+   "lr": 0.001709,
+   "tokens_per_second": 921305,
+   "minutes": 27.22,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 5600,
+   "train_loss": 2.9965,
+   "lr": 0.001691,
+   "tokens_per_second": 921367,
+   "minutes": 27.46,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 5650,
+   "train_loss": 3.0261,
+   "lr": 0.001672,
+   "tokens_per_second": 921494,
+   "minutes": 27.7,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 5700,
+   "train_loss": 2.9857,
+   "lr": 0.001654,
+   "tokens_per_second": 921931,
+   "minutes": 27.93,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 5750,
+   "train_loss": 2.9916,
+   "lr": 0.001635,
+   "tokens_per_second": 921797,
+   "minutes": 28.17,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 5800,
+   "train_loss": 2.9868,
+   "lr": 0.001616,
+   "tokens_per_second": 921701,
+   "minutes": 28.41,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 5850,
+   "train_loss": 3.0135,
+   "lr": 0.001598,
+   "tokens_per_second": 921157,
+   "minutes": 28.64,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 5900,
+   "train_loss": 3.0171,
+   "lr": 0.001579,
+   "tokens_per_second": 921347,
+   "minutes": 28.88,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 5950,
+   "train_loss": 3.0381,
+   "lr": 0.001561,
+   "tokens_per_second": 921107,
+   "minutes": 29.12,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 6000,
+   "train_loss": 3.0234,
+   "lr": 0.001542,
+   "tokens_per_second": 920618,
+   "minutes": 29.35,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1,
+   "val_loss": 3.014
+  },
+  {
+   "step": 6050,
+   "train_loss": 2.9671,
+   "lr": 0.001524,
+   "tokens_per_second": 920459,
+   "minutes": 29.64,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 6100,
+   "train_loss": 3.0432,
+   "lr": 0.001505,
+   "tokens_per_second": 919920,
+   "minutes": 29.88,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 6150,
+   "train_loss": 3.0124,
+   "lr": 0.001487,
+   "tokens_per_second": 920638,
+   "minutes": 30.11,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 6200,
+   "train_loss": 3.0045,
+   "lr": 0.001468,
+   "tokens_per_second": 919978,
+   "minutes": 30.35,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 6250,
+   "train_loss": 3.0096,
+   "lr": 0.00145,
+   "tokens_per_second": 919568,
+   "minutes": 30.59,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 6300,
+   "train_loss": 2.9723,
+   "lr": 0.001432,
+   "tokens_per_second": 919292,
+   "minutes": 30.83,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 6350,
+   "train_loss": 2.9286,
+   "lr": 0.001413,
+   "tokens_per_second": 919775,
+   "minutes": 31.06,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 6400,
+   "train_loss": 3.0349,
+   "lr": 0.001395,
+   "tokens_per_second": 920505,
+   "minutes": 31.3,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 6450,
+   "train_loss": 2.9374,
+   "lr": 0.001377,
+   "tokens_per_second": 920417,
+   "minutes": 31.54,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 6500,
+   "train_loss": 2.9625,
+   "lr": 0.001359,
+   "tokens_per_second": 919773,
+   "minutes": 31.78,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1,
+   "val_loss": 2.9999
+  },
+  {
+   "step": 6550,
+   "train_loss": 2.9994,
+   "lr": 0.001341,
+   "tokens_per_second": 919974,
+   "minutes": 32.06,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 6600,
+   "train_loss": 2.9899,
+   "lr": 0.001323,
+   "tokens_per_second": 920286,
+   "minutes": 32.3,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 6650,
+   "train_loss": 2.9663,
+   "lr": 0.001305,
+   "tokens_per_second": 920398,
+   "minutes": 32.54,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 6700,
+   "train_loss": 2.9463,
+   "lr": 0.001287,
+   "tokens_per_second": 921337,
+   "minutes": 32.77,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 6750,
+   "train_loss": 2.9973,
+   "lr": 0.001269,
+   "tokens_per_second": 921370,
+   "minutes": 33.01,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 6800,
+   "train_loss": 2.9206,
+   "lr": 0.001251,
+   "tokens_per_second": 920742,
+   "minutes": 33.25,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 6850,
+   "train_loss": 2.9574,
+   "lr": 0.001233,
+   "tokens_per_second": 920284,
+   "minutes": 33.48,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 6900,
+   "train_loss": 2.9734,
+   "lr": 0.001216,
+   "tokens_per_second": 919909,
+   "minutes": 33.72,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 6950,
+   "train_loss": 2.992,
+   "lr": 0.001198,
+   "tokens_per_second": 919380,
+   "minutes": 33.96,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 7000,
+   "train_loss": 2.9482,
+   "lr": 0.001181,
+   "tokens_per_second": 919305,
+   "minutes": 34.2,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1,
+   "val_loss": 2.9834
+  },
+  {
+   "step": 7050,
+   "train_loss": 2.9944,
+   "lr": 0.001163,
+   "tokens_per_second": 919402,
+   "minutes": 34.48,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 7100,
+   "train_loss": 2.9649,
+   "lr": 0.001146,
+   "tokens_per_second": 919490,
+   "minutes": 34.72,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 7150,
+   "train_loss": 2.9533,
+   "lr": 0.001129,
+   "tokens_per_second": 919450,
+   "minutes": 34.96,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 7200,
+   "train_loss": 2.9467,
+   "lr": 0.001112,
+   "tokens_per_second": 919809,
+   "minutes": 35.19,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 7250,
+   "train_loss": 2.9482,
+   "lr": 0.001095,
+   "tokens_per_second": 920118,
+   "minutes": 35.43,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 7300,
+   "train_loss": 2.9288,
+   "lr": 0.001078,
+   "tokens_per_second": 920161,
+   "minutes": 35.67,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 7350,
+   "train_loss": 2.9711,
+   "lr": 0.001061,
+   "tokens_per_second": 920166,
+   "minutes": 35.91,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 7400,
+   "train_loss": 2.9711,
+   "lr": 0.001045,
+   "tokens_per_second": 920031,
+   "minutes": 36.14,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 7450,
+   "train_loss": 2.9619,
+   "lr": 0.001028,
+   "tokens_per_second": 920234,
+   "minutes": 36.38,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 7500,
+   "train_loss": 2.9199,
+   "lr": 0.001012,
+   "tokens_per_second": 920421,
+   "minutes": 36.62,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1,
+   "val_loss": 2.9683
+  },
+  {
+   "step": 7550,
+   "train_loss": 2.9091,
+   "lr": 0.000995,
+   "tokens_per_second": 921260,
+   "minutes": 36.9,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 7600,
+   "train_loss": 2.9747,
+   "lr": 0.000979,
+   "tokens_per_second": 920906,
+   "minutes": 37.14,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 7650,
+   "train_loss": 2.9445,
+   "lr": 0.000963,
+   "tokens_per_second": 920371,
+   "minutes": 37.38,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 7700,
+   "train_loss": 2.9593,
+   "lr": 0.000947,
+   "tokens_per_second": 920018,
+   "minutes": 37.62,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 7750,
+   "train_loss": 2.9826,
+   "lr": 0.000931,
+   "tokens_per_second": 920997,
+   "minutes": 37.85,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 7800,
+   "train_loss": 2.933,
+   "lr": 0.000916,
+   "tokens_per_second": 921462,
+   "minutes": 38.09,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 7850,
+   "train_loss": 2.9202,
+   "lr": 0.0009,
+   "tokens_per_second": 921036,
+   "minutes": 38.33,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 7900,
+   "train_loss": 2.9152,
+   "lr": 0.000885,
+   "tokens_per_second": 920538,
+   "minutes": 38.56,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 7950,
+   "train_loss": 2.926,
+   "lr": 0.00087,
+   "tokens_per_second": 919615,
+   "minutes": 38.8,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 8000,
+   "train_loss": 2.8966,
+   "lr": 0.000855,
+   "tokens_per_second": 919288,
+   "minutes": 39.04,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1,
+   "val_loss": 2.9544
+  },
+  {
+   "step": 8050,
+   "train_loss": 2.9517,
+   "lr": 0.00084,
+   "tokens_per_second": 919884,
+   "minutes": 39.32,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 8100,
+   "train_loss": 2.9257,
+   "lr": 0.000825,
+   "tokens_per_second": 921250,
+   "minutes": 39.56,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 8150,
+   "train_loss": 2.9299,
+   "lr": 0.00081,
+   "tokens_per_second": 921425,
+   "minutes": 39.8,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 8200,
+   "train_loss": 2.934,
+   "lr": 0.000796,
+   "tokens_per_second": 921296,
+   "minutes": 40.04,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 8250,
+   "train_loss": 2.8999,
+   "lr": 0.000782,
+   "tokens_per_second": 920909,
+   "minutes": 40.27,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 8300,
+   "train_loss": 2.9048,
+   "lr": 0.000767,
+   "tokens_per_second": 920173,
+   "minutes": 40.51,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 8350,
+   "train_loss": 2.8714,
+   "lr": 0.000753,
+   "tokens_per_second": 919853,
+   "minutes": 40.75,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 8400,
+   "train_loss": 2.8722,
+   "lr": 0.00074,
+   "tokens_per_second": 919551,
+   "minutes": 40.98,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 8450,
+   "train_loss": 2.9693,
+   "lr": 0.000726,
+   "tokens_per_second": 920679,
+   "minutes": 41.22,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 8500,
+   "train_loss": 2.9393,
+   "lr": 0.000713,
+   "tokens_per_second": 921776,
+   "minutes": 41.46,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1,
+   "val_loss": 2.9392
+  },
+  {
+   "step": 8550,
+   "train_loss": 2.8964,
+   "lr": 0.000699,
+   "tokens_per_second": 920913,
+   "minutes": 41.74,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 8600,
+   "train_loss": 2.9192,
+   "lr": 0.000686,
+   "tokens_per_second": 919992,
+   "minutes": 41.98,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 8650,
+   "train_loss": 2.9288,
+   "lr": 0.000673,
+   "tokens_per_second": 919914,
+   "minutes": 42.22,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 8700,
+   "train_loss": 2.8246,
+   "lr": 0.000661,
+   "tokens_per_second": 919626,
+   "minutes": 42.46,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 8750,
+   "train_loss": 2.9473,
+   "lr": 0.000648,
+   "tokens_per_second": 919635,
+   "minutes": 42.69,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 8800,
+   "train_loss": 2.8781,
+   "lr": 0.000636,
+   "tokens_per_second": 921027,
+   "minutes": 42.93,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 8850,
+   "train_loss": 2.8393,
+   "lr": 0.000624,
+   "tokens_per_second": 921693,
+   "minutes": 43.17,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 8900,
+   "train_loss": 2.9105,
+   "lr": 0.000612,
+   "tokens_per_second": 922213,
+   "minutes": 43.4,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 8950,
+   "train_loss": 2.9286,
+   "lr": 0.0006,
+   "tokens_per_second": 923184,
+   "minutes": 43.64,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 9000,
+   "train_loss": 2.8668,
+   "lr": 0.000588,
+   "tokens_per_second": 923153,
+   "minutes": 43.88,
+   "mfu_percent": 22.3,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1,
+   "val_loss": 2.9286
+  },
+  {
+   "step": 9050,
+   "train_loss": 2.9063,
+   "lr": 0.000577,
+   "tokens_per_second": 922088,
+   "minutes": 44.16,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 9100,
+   "train_loss": 2.9428,
+   "lr": 0.000566,
+   "tokens_per_second": 921296,
+   "minutes": 44.4,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 9150,
+   "train_loss": 2.9146,
+   "lr": 0.000555,
+   "tokens_per_second": 920478,
+   "minutes": 44.64,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 9200,
+   "train_loss": 2.9337,
+   "lr": 0.000544,
+   "tokens_per_second": 919857,
+   "minutes": 44.87,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 9250,
+   "train_loss": 2.8898,
+   "lr": 0.000533,
+   "tokens_per_second": 919677,
+   "minutes": 45.11,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 9300,
+   "train_loss": 2.8907,
+   "lr": 0.000523,
+   "tokens_per_second": 919018,
+   "minutes": 45.35,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 9350,
+   "train_loss": 2.9104,
+   "lr": 0.000513,
+   "tokens_per_second": 919151,
+   "minutes": 45.59,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 9400,
+   "train_loss": 2.9346,
+   "lr": 0.000503,
+   "tokens_per_second": 919820,
+   "minutes": 45.82,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 9450,
+   "train_loss": 2.864,
+   "lr": 0.000493,
+   "tokens_per_second": 920435,
+   "minutes": 46.06,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 9500,
+   "train_loss": 2.8496,
+   "lr": 0.000484,
+   "tokens_per_second": 920689,
+   "minutes": 46.3,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1,
+   "val_loss": 2.9166
+  },
+  {
+   "step": 9550,
+   "train_loss": 2.9794,
+   "lr": 0.000475,
+   "tokens_per_second": 920761,
+   "minutes": 46.58,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 9600,
+   "train_loss": 2.8557,
+   "lr": 0.000466,
+   "tokens_per_second": 921094,
+   "minutes": 46.82,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 9650,
+   "train_loss": 2.8992,
+   "lr": 0.000457,
+   "tokens_per_second": 921528,
+   "minutes": 47.06,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 9700,
+   "train_loss": 2.835,
+   "lr": 0.000448,
+   "tokens_per_second": 920716,
+   "minutes": 47.29,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 9750,
+   "train_loss": 2.9188,
+   "lr": 0.00044,
+   "tokens_per_second": 920417,
+   "minutes": 47.53,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 9800,
+   "train_loss": 2.925,
+   "lr": 0.000432,
+   "tokens_per_second": 920052,
+   "minutes": 47.77,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 9850,
+   "train_loss": 2.9457,
+   "lr": 0.000424,
+   "tokens_per_second": 919561,
+   "minutes": 48.01,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 9900,
+   "train_loss": 2.9236,
+   "lr": 0.000416,
+   "tokens_per_second": 919301,
+   "minutes": 48.24,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 9950,
+   "train_loss": 2.8542,
+   "lr": 0.000409,
+   "tokens_per_second": 919176,
+   "minutes": 48.48,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 10000,
+   "train_loss": 2.8907,
+   "lr": 0.000402,
+   "tokens_per_second": 919788,
+   "minutes": 48.72,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1,
+   "val_loss": 2.9086
+  },
+  {
+   "step": 10050,
+   "train_loss": 2.7981,
+   "lr": 0.000395,
+   "tokens_per_second": 920621,
+   "minutes": 49.0,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 10100,
+   "train_loss": 2.8218,
+   "lr": 0.000388,
+   "tokens_per_second": 920907,
+   "minutes": 49.24,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 10150,
+   "train_loss": 2.8762,
+   "lr": 0.000382,
+   "tokens_per_second": 921101,
+   "minutes": 49.48,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 10200,
+   "train_loss": 2.8884,
+   "lr": 0.000375,
+   "tokens_per_second": 920843,
+   "minutes": 49.72,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 10250,
+   "train_loss": 2.8617,
+   "lr": 0.000369,
+   "tokens_per_second": 921002,
+   "minutes": 49.95,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 10300,
+   "train_loss": 2.8414,
+   "lr": 0.000364,
+   "tokens_per_second": 921151,
+   "minutes": 50.19,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 10350,
+   "train_loss": 2.8439,
+   "lr": 0.000358,
+   "tokens_per_second": 920309,
+   "minutes": 50.43,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 10400,
+   "train_loss": 2.8992,
+   "lr": 0.000353,
+   "tokens_per_second": 919693,
+   "minutes": 50.66,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 10450,
+   "train_loss": 2.8711,
+   "lr": 0.000348,
+   "tokens_per_second": 919270,
+   "minutes": 50.9,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 10500,
+   "train_loss": 2.8533,
+   "lr": 0.000343,
+   "tokens_per_second": 920712,
+   "minutes": 51.14,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1,
+   "val_loss": 2.9005
+  },
+  {
+   "step": 10550,
+   "train_loss": 2.8638,
+   "lr": 0.000338,
+   "tokens_per_second": 921492,
+   "minutes": 51.42,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 10600,
+   "train_loss": 2.8673,
+   "lr": 0.000334,
+   "tokens_per_second": 921466,
+   "minutes": 51.66,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 10650,
+   "train_loss": 2.8849,
+   "lr": 0.00033,
+   "tokens_per_second": 920910,
+   "minutes": 51.9,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 10700,
+   "train_loss": 2.8652,
+   "lr": 0.000326,
+   "tokens_per_second": 920277,
+   "minutes": 52.14,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 10750,
+   "train_loss": 2.8523,
+   "lr": 0.000323,
+   "tokens_per_second": 919903,
+   "minutes": 52.37,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 10800,
+   "train_loss": 2.8797,
+   "lr": 0.00032,
+   "tokens_per_second": 919671,
+   "minutes": 52.61,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 10850,
+   "train_loss": 2.858,
+   "lr": 0.000317,
+   "tokens_per_second": 919122,
+   "minutes": 52.85,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 10900,
+   "train_loss": 2.8607,
+   "lr": 0.000314,
+   "tokens_per_second": 919161,
+   "minutes": 53.09,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 10950,
+   "train_loss": 2.8733,
+   "lr": 0.000311,
+   "tokens_per_second": 920086,
+   "minutes": 53.32,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 11000,
+   "train_loss": 2.8539,
+   "lr": 0.000309,
+   "tokens_per_second": 920688,
+   "minutes": 53.56,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1,
+   "val_loss": 2.8962
+  },
+  {
+   "step": 11050,
+   "train_loss": 2.8992,
+   "lr": 0.000307,
+   "tokens_per_second": 920928,
+   "minutes": 53.84,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 11100,
+   "train_loss": 2.8486,
+   "lr": 0.000305,
+   "tokens_per_second": 920127,
+   "minutes": 54.08,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 11150,
+   "train_loss": 2.8515,
+   "lr": 0.000304,
+   "tokens_per_second": 919904,
+   "minutes": 54.32,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 11200,
+   "train_loss": 2.8647,
+   "lr": 0.000302,
+   "tokens_per_second": 919524,
+   "minutes": 54.56,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 11250,
+   "train_loss": 2.8809,
+   "lr": 0.000301,
+   "tokens_per_second": 920127,
+   "minutes": 54.79,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 11300,
+   "train_loss": 2.8752,
+   "lr": 0.000301,
+   "tokens_per_second": 921273,
+   "minutes": 55.03,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 11350,
+   "train_loss": 2.9045,
+   "lr": 0.0003,
+   "tokens_per_second": 920772,
+   "minutes": 55.27,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 11400,
+   "train_loss": 2.7934,
+   "lr": 0.0003,
+   "tokens_per_second": 920858,
+   "minutes": 55.51,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1
+  },
+  {
+   "step": 11420,
+   "train_loss": 2.8882,
+   "lr": 0.0003,
+   "tokens_per_second": 921250,
+   "minutes": 55.6,
+   "mfu_percent": 22.2,
+   "gpu_util_percent": 100,
+   "gpu_memory_gb": 85.1,
+   "val_loss": 2.8934
+  }
+ ],
+ "exports": {
+  "parity_prompts": 300,
+  "parity_steps_per_prompt": 5,
+  "float": {
+   "megabytes": 60.81,
+   "same_first_token": 1.0,
+   "same_five_greedy_tokens": 1.0,
+   "cached_step_ms_p50_cpu": 5.544
+  },
+  "int8": {
+   "megabytes": 15.64,
+   "same_first_token": 0.0,
+   "same_five_greedy_tokens": 0.0,
+   "cached_step_ms_p50_cpu": 4.808
+  }
+ }
+}

v2/experiments/tinynext-v2-3b/train.py

@@ -0,0 +1,691 @@
+# /// script
+# requires-python = ">=3.11,<3.14"
+# dependencies = [
+#   "torch>=2.6,<2.8",
+#   "datasets>=3.0",
+#   "numpy<3",
+#   "pyarrow>=15",
+#   "tokenizers>=0.20",
+#   "transformers>=4.46",
+#   "onnx>=1.17",
+#   "onnxruntime>=1.20",
+#   "safetensors>=0.4",
+#   "huggingface_hub>=0.30",
+#   "nvidia-ml-py>=12",
+#   "trackio>=0.2",
+# ]
+# ///
+"""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
+
+
+# ---- Shared eval code ------------------------------------------------------------------------------------------
+
+
+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
+
+
+# ---- Data ------------------------------------------------------------------------------------------------------
+
+
+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}
+
+
+# Byte-level BPE with an 8K vocabulary averages a little over 4 characters per token on English. Reading 40% more
+# text than that covers the variance, and anything short gets repeated (and logged) in flatten().
+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 = []  # also learn the plain "nothing after the cursor" case in this format
+    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
+
+
+# ---- Model -----------------------------------------------------------------------------------------------------
+
+
+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  # set only for export: an untied copy of the output weights, so int8 quantization covers it
+        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
+
+
+# ---- Training --------------------------------------------------------------------------------------------------
+
+
+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:  # no NVIDIA GPU (a laptop smoke run): report nothing rather than fail
+            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()  # syncs the GPU, so throughput below is real
+            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),
+                # Training compute is about 6 FLOPs per parameter per token; an A100 does 312 TFLOP/s in bf16.
+                "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
+
+
+# ---- Eval, export --------------------------------------------------------------------------------------------
+
+
+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]
+    # A model trained without fill-in-the-middle never saw the special tokens, so it only gets the text before.
+    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)
+    # Trace with a non-empty cache so no shape gets fixed at zero; the checks below start from an empty one.
+    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
+
+
+# ---- Prepared data ---------------------------------------------------------------------------------------------
+
+
+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"]
+
+
+# ---- Main ------------------------------------------------------------------------------------------------------
+
+
+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
+
+    # The whole token array lives on the GPU; each step just indexes into it.
+    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()