remove the 3b-token experiment folder, now the same as v2

README.md

@@ -157,7 +157,7 @@ The browser computes in float32 on purpose: onnxruntime-web's int8 kernels were
 | `v2/eval.py`, `v2/eval_set.json` | The eval harness and the fixed text it scores |
 | `v2/export.py`, `v2/compress.py` | Re-export to ONNX, and the int8 download format |
 | `v2/metrics/*.json` | Eval results for v1, v2, the browser format, SmolLM2-135M and Pythia 14M–160M |
-| `v2/experiments/` | The no-FIM run, and the 3B-token run as trained (the same weights as `v2/`) |
+| `v2/experiments/` | The no-FIM run, for comparison |
 | `v2/report.json` | Training settings, loss history and export checks |
 
 </details>

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

@@ -1,16 +0,0 @@
-{
- "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

@@ -1,499 +0,0 @@
-# /// 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

@@ -1,371 +0,0 @@
-{
- "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

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

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

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

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

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

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

@@ -1,2390 +0,0 @@
-{
- "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

@@ -1,691 +0,0 @@
-# /// 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()