parallel data prep, data on GPU, trackio, safetensors

v2/train.py

@@ -4,21 +4,29 @@
 #   "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/.
+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 --flavor a100-large --secrets HF_TOKEN train.py
-  uv run train.py --smoke --no-push      # a few minutes on a laptop, to check the whole pipeline
+  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
@@ -28,6 +36,7 @@ import math
 import os
 import random
 import shutil
+import sys
 import time
 from pathlib import Path
 
@@ -40,12 +49,21 @@ 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."""
+    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
@@ -54,6 +72,7 @@ def load_eval_module(repo: str, folder: str):
             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
 
@@ -61,80 +80,69 @@ def load_eval_module(repo: str, folder: str):
 # ---- Data ------------------------------------------------------------------------------------------------------
 
 
-def fineweb_docs():
-    from datasets import load_dataset
-
-    for row in load_dataset("HuggingFaceFW/fineweb-edu", name="sample-10BT", split="train", streaming=True):
-        yield row["text"]
+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_docs():
-    from datasets import load_dataset
 
-    for row in load_dataset("HuggingFaceTB/smoltalk", "all", split="train", streaming=True):
-        turns = [m["content"] for m in row["messages"] if m["role"] in ("user", "assistant")]
-        if turns:
-            yield "\n\n".join(turns)
+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_docs(ev):
-    """Training senders only: the eval's held-out senders (and enron-autocomplete's val split) are skipped."""
+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}
+
+
+def load_sources(args) -> dict:
+    """Downloads whole parquet files and cleans them on every core."""
     from datasets import load_dataset
 
-    for row in load_dataset("corbt/enron-emails", split="train", streaming=True):
-        sender = (row["from"] or "").lower()
-        if "@enron.com" not in sender or ev.BULK_SENDER.search(sender):
-            continue
-        if int(hashlib.sha1(sender.encode()).hexdigest(), 16) % 100 < 10:
-            continue
-        paragraphs = ev.clean_email(row["body"] or "")
-        if paragraphs:
-            yield "\n\n".join(paragraphs)
-
-
-def mixed_docs(ev, weights: dict[str, float], seed: int):
-    """Interleaves the sources by weight, forever (until the caller stops)."""
-    rng = random.Random(seed)
-    sources = {"fineweb": fineweb_docs(), "smoltalk": smoltalk_docs(), "enron": enron_docs(ev)}
-    names, probs = zip(*weights.items())
-    counts = dict.fromkeys(names, 0)
-    while True:
-        name = rng.choices(names, probs)[0]
-        try:
-            doc = next(sources[name])
-        except StopIteration:
-            sources[name] = {"fineweb": fineweb_docs, "smoltalk": smoltalk_docs}.get(name, lambda: enron_docs(ev))()
-            continue
-        counts[name] += 1
-        # Log the first document from each source and then every 50k, so a stalled download shows which one.
-        if counts[name] == 1 or counts[name] % 50_000 == 0:
-            print(f"  data seed {seed}: {name} doc {counts[name]:,}", flush=True)
-        yield doc
-
-
-def clean(text: str, ev) -> str:
-    """The browser model only sees keyboard characters, so train on the same."""
-    return text.translate(ev.SMART_CHARS).encode("ascii", "ignore").decode().strip()
+    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"]).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")
+    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(docs, vocab: int, sample_docs: int):
+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((next(docs) for _ in range(sample_docs)), trainer=trainer)
+    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
+        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]
@@ -143,30 +151,58 @@ def fim(ids: list[int], rng: random.Random, max_middle: int, max_suffix: int, ma
     return [SUF, *suffix, PRE, *prefix, MID, *middle, EOT]
 
 
-def build_tokens(tokenizer, docs, ev, total: int, seq: int, fim_rate: float, seed: int) -> np.ndarray:
-    """Packs documents into rows of seq+1 tokens. A share of chunks become fill-in-the-middle examples."""
-    rng = random.Random(seed)
-    out = np.empty(total + seq + 1, dtype=np.uint16)
-    filled, started, next_report = 0, time.time(), 0
-    while filled < total:
-        batch = [clean(next(docs), ev) for _ in range(512)]
-        for ids in tokenizer.encode_batch([b for b in batch if b]):
-            ids = ids.ids
+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]
-                piece = fim(chunk, rng, 24, 64, seq - 100) if rng.random() < fim_rate else chunk + [EOT]
-                take = min(len(piece), len(out) - filled)
-                out[filled: filled + take] = piece[:take]
-                filled += take
-                if filled >= total:
-                    break
-            if filled >= total:
-                break
-        if filled >= next_report:
-            print(f"  tokens {filled / 1e6:.0f}M / {total / 1e6:.0f}M ({time.time() - started:.0f}s)", flush=True)
-            next_report += 20_000_000
-    rows = total // (seq + 1)
-    return out[: rows * (seq + 1)].reshape(rows, seq + 1)
+                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 -----------------------------------------------------------------------------------------------------
@@ -261,32 +297,57 @@ class BrowserStep(nn.Module):
 # ---- Training --------------------------------------------------------------------------------------------------
 
 
-def evaluate_loss(model, rows: np.ndarray, device, batch: int) -> float:
+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 = torch.from_numpy(rows[at: at + batch].astype(np.int64)).to(device)
+            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: np.ndarray, val_rows: np.ndarray, args, device) -> list[dict]:
+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 n, p in model.named_parameters() if p.dim() >= 2]
-    no_decay = [p for n, p in model.named_parameters() if p.dim() < 2]
+    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
-    history, rng, started = [], np.random.default_rng(0), time.time()
+    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 = rng.integers(0, len(rows), args.batch)
-        chunk = torch.from_numpy(rows[index].astype(np.int64)).to(device, non_blocking=True)
+        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())
@@ -295,12 +356,23 @@ def train(model, rows: np.ndarray, val_rows: np.ndarray, args, device) -> list[d
         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:
-            entry = {"step": step, "train_loss": round(loss.item(), 4), "minutes": round((time.time() - started) / 60, 2)}
+            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
 
 
@@ -365,6 +437,18 @@ def export(model: TinyNext, output: Path, device) -> dict:
     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
+
+
 # ---- Main ------------------------------------------------------------------------------------------------------
 
 
@@ -375,58 +459,91 @@ def parse_args():
     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=256)
-    parser.add_argument("--tokens", type=float, default=800e6, help="Training tokens to build")
+    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=256)
+    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=100)
+    parser.add_argument("--log-every", type=int, default=50)
     parser.add_argument("--repo", default="cfahlgren1/tinynext")
     parser.add_argument("--folder", default="v2")
     parser.add_argument("--name", default="tinynext-v2")
+    parser.add_argument("--trackio-space", default="cfahlgren1/tinynext-trackio")
     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 = 2048, 64, 2, 2, 128
-        args.tokens, args.batch, args.tokenizer_docs, args.eval_pauses, args.log_every = 3e6, 32, 3000, 20, 20
+        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.folder)
     weights = {k: float(v) for k, v in (part.split("=") for part in args.mix.split(","))}
-    print(f"device {device}, mix {weights}", flush=True)
+    print(f"device {device}, {cpus()} CPUs, mix {weights}", flush=True)
+    tracker = start_tracker(args)
 
-    tokenizer = train_tokenizer((clean(d, ev) for d in mixed_docs(ev, weights, seed=1)), args.vocab, args.tokenizer_docs)
+    started = time.time()
+    sources = load_sources(args)
+    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"))
-    # Validation takes the first documents of the stream and training continues after them, so they never overlap.
-    docs = mixed_docs(ev, weights, seed=2)
-    val_rows = build_tokens(tokenizer, docs, ev, int(min(2e6, args.tokens / 20)), args.seq, args.fim_rate, seed=5)
-    rows = build_tokens(tokenizer, docs, ev, int(args.tokens), args.seq, args.fim_rate, seed=3)
+    rows, val_rows, counts = build_rows(sources, tokenizer, weights, args)
+    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)
 
+    # 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, {rows.size / 1e6:.0f}M tokens", flush=True)
+    print(f"{parameters / 1e6:.2f}M parameters", flush=True)
     started = time.time()
-    history = train(model, rows, val_rows, args, device)
-    minutes = (time.time() - started) / 60
+    history = train(model, rows_gpu, val_gpu, args, device, tracker)
+    train_minutes = (time.time() - started) / 60
 
     model.eval()
-    torch.save({"config": model.config, "state_dict": model.state_dict()}, output / "model.pt")
+    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), 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)
 
     config = {
@@ -434,16 +551,19 @@ def main():
         "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, "train_minutes": round(minutes, 1), "history": history, "exports": exported}
+    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"):
-            shutil.copy(path, output / "metrics" / path.name)
+            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
@@ -451,7 +571,7 @@ def main():
 
     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, trained {minutes:.0f} min",
+        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}")