@tuned-tensor/cli 0.13.1-beta.0 → 0.13.1-beta.2

package/docs/local-runtime/architecture.md

@@ -40,6 +40,30 @@ The first version intentionally excludes:
 The next model or method should arrive with its own locked dependencies,
 resource defaults, data contract, and real CUDA acceptance test.
 
+TT Local currently ships two such methods:
+
+- **Adapter** (`training/local-runner`): certified Transformers checkpoints and
+  PEFT LoRA SFT. Driven by `tt run` and adapter pipelines.
+- **Foundation** (`training/foundation`): a from-scratch BPE tokenizer and GPT
+  with no Transformers/PEFT. Driven only by `tt pipeline` from a foundation
+  `tunedtensor.json`. Pretrain, SFT, and RL require CUDA. Tokenizer training
+  and the CPU unit tests do not. The current runtime is single-process and
+  rejects `nproc_per_node` values above 1 and unsupported steps such as
+  `compare` before creating run artifacts.
+
+The foundation engine is a readable lifecycle baseline inspired by nanochat,
+not a port of nanochat's data, distributed, optimizer, or compute-optimal
+training stack. Its bounded pretraining corpus repeats the foundation spec's
+prompt and examples, and chat evaluation reuses those examples. RL uses a
+single seeded on-policy sample per step with a sparse last-number exact reward
+and requires every expected output to contain a number. BPB is normalized over
+decoded target bytes from each evaluated window, but the corpus remains
+example-derived. The resulting metrics prove stage wiring and can demonstrate
+deliberate overfitting; they do not measure held-out model capability.
+
+Do not mix those lockfiles. Each engine hashes its own `pyproject.toml` and
+`uv.lock` into a distinct cache environment.
+
 ## Python boundary
 
 Node owns input validation, state, process lifecycle, manifests, and reports.

package/examples/local-runtime/foundation-spec.json

@@ -0,0 +1,31 @@
+{
+  "engine": "foundation",
+  "name": "Foundation smoke chat",
+  "description": "Tiny from-scratch tokenizer, pretrain, chat SFT, and eval.",
+  "system_prompt": "You are a helpful assistant.",
+  "guidelines": [
+    "Answer the user request directly."
+  ],
+  "constraints": [],
+  "examples": [
+    {
+      "input": "Say hello in one sentence.",
+      "output": "Hello — how can I help you today?"
+    },
+    {
+      "input": "What is 2 + 2?",
+      "output": "2 + 2 = 4."
+    }
+  ],
+  "foundation": {
+    "vocab_size": 256,
+    "max_chars": 20000,
+    "depth": 2,
+    "pretrain_steps": 2,
+    "finetune_steps": 2,
+    "rl_steps": 0,
+    "batch_size": 2,
+    "sequence_length": 64,
+    "nproc_per_node": 1
+  }
+}

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@tuned-tensor/cli",
-  "version": "0.13.1-beta.0",
+  "version": "0.13.1-beta.2",
   "description": "Local TT agent and CUDA fine-tuning CLI for Tuned Tensor workflows",
   "repository": {
     "type": "git",
@@ -17,7 +17,7 @@
     "test": "npm run test:cli && npm run test:local-runtime && npm run test:python",
     "test:cli": "vitest run",
     "test:local-runtime": "node --import tsx --test \"test/local-runtime/**/*.test.ts\"",
-    "test:python": "python3 -m unittest discover -s training/local-runner/tests -p 'test_*.py'",
+    "test:python": "python3 -m unittest discover -s training/local-runner/tests -p 'test_*.py' && uv run --frozen --project training/foundation python -m unittest discover -s training/foundation/tests -p 'test_*.py'",
     "test:watch": "vitest"
   },
   "files": [
@@ -25,6 +25,9 @@
     "training/local-runner/pyproject.toml",
     "training/local-runner/uv.lock",
     "training/local-runner/src/*.py",
+    "training/foundation/pyproject.toml",
+    "training/foundation/uv.lock",
+    "training/foundation/src/*.py",
     "docs/local-runtime",
     "examples/local-runtime/*.json",
     "examples/local-runtime/*.jsonl"
@@ -44,7 +47,7 @@
   "dependencies": {
     "@earendil-works/pi-agent-core": "0.84.1",
     "@earendil-works/pi-coding-agent": "0.84.1",
-    "@tuned-tensor/pipeline-contract": "1.0.0",
+    "@tuned-tensor/pipeline-contract": "1.1.0",
     "chalk": "^5.6.2",
     "cli-table3": "^0.6.5",
     "commander": "^14.0.3",

package/training/foundation/pyproject.toml

@@ -0,0 +1,24 @@
+[project]
+name = "tuned-tensor-foundation"
+version = "0.1.0"
+requires-python = ">=3.11"
+dependencies = [
+  "huggingface-hub>=0.36.0",
+  "numpy>=2.0.0",
+  "safetensors>=0.4.5",
+  "tokenizers>=0.21.0",
+  "torch>=2.11.0",
+]
+
+[tool.uv.sources]
+torch = [
+  { index = "pytorch-cu130", marker = "sys_platform == 'linux'" },
+]
+
+[[tool.uv.index]]
+name = "pytorch-cu130"
+url = "https://download.pytorch.org/whl/cu130"
+explicit = true
+
+[tool.uv]
+package = false

package/training/foundation/src/common.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from typing import Any
+
+
+def ensure_private_directory(path: str | Path) -> Path:
+    destination = Path(path)
+    destination.mkdir(parents=True, exist_ok=True, mode=0o700)
+    destination.chmod(0o700)
+    return destination
+
+
+def make_private_file(path: str | Path) -> Path:
+    destination = Path(path)
+    destination.chmod(0o600)
+    return destination
+
+
+def load_config(argv: list[str] | None = None) -> dict[str, Any]:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", required=True)
+    args = parser.parse_args(argv)
+    return json.loads(Path(args.config).read_text(encoding="utf-8"))
+
+
+def write_json(path: str | Path, value: Any) -> None:
+    destination = Path(path)
+    ensure_private_directory(destination.parent)
+    with destination.open("w", encoding="utf-8") as output:
+        make_private_file(destination)
+        output.write(json.dumps(value, indent=2) + "\n")
+
+
+def require_cuda(action: str) -> None:
+    import torch
+
+    if not torch.cuda.is_available():
+        raise SystemExit(f"Foundation {action} requires CUDA.")

package/training/foundation/src/data.py

@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+import hashlib
+import re
+from typing import Any
+
+from tokenizers import Tokenizer
+from tokenizers import models as tokenizer_models
+from tokenizers import pre_tokenizers, trainers
+
+SYSTEM = "<|system|>"
+USER = "<|user|>"
+ASSISTANT = "<|assistant|>"
+END = "<|end|>"
+UNK = "<unk>"
+PAD = "<pad>"
+SPECIAL_TOKENS = [UNK, PAD, SYSTEM, USER, ASSISTANT, END]
+IGNORE_INDEX = -100
+
+
+def train_tokenizer(corpus: str, vocab_size: int) -> Tokenizer:
+    tokenizer = Tokenizer(tokenizer_models.BPE(unk_token=UNK))
+    tokenizer.pre_tokenizer = pre_tokenizers.Whitespace()
+    trainer = trainers.BpeTrainer(
+        vocab_size=max(int(vocab_size), len(SPECIAL_TOKENS) + 32),
+        special_tokens=SPECIAL_TOKENS,
+        min_frequency=1,
+    )
+    tokenizer.train_from_iterator([corpus], trainer=trainer)
+    return tokenizer
+
+
+def smoke_corpus(examples: list[dict[str, str]], max_chars: int, extra: str = "") -> str:
+    chunks = [extra, *[example.get("input", "") for example in examples], *[example.get("output", "") for example in examples]]
+    seed = "\n".join(part for part in chunks if part).strip() or "the quick brown fox jumps over the lazy dog"
+    seed = f"{seed}\n"
+    repeats = max(1, (max_chars // max(len(seed), 1)) + 1)
+    return (seed * repeats)[:max_chars]
+
+
+def corpus_hash(text: str) -> str:
+    return hashlib.sha256(text.encode("utf-8")).hexdigest()
+
+
+def format_prompt(system_prompt: str, user: str) -> str:
+    system = system_prompt.strip() or "You are a helpful assistant."
+    return f"{SYSTEM}{system}{END}{USER}{user}{END}{ASSISTANT}"
+
+
+def format_chat(system_prompt: str, user: str, assistant: str) -> str:
+    return f"{format_prompt(system_prompt, user)}{assistant}{END}"
+
+
+def encode_ids(tokenizer: Tokenizer, text: str) -> list[int]:
+    return list(tokenizer.encode(text).ids)
+
+
+def decode_ids(tokenizer: Tokenizer, ids: list[int]) -> str:
+    return tokenizer.decode(ids, skip_special_tokens=False)
+
+
+def decoded_byte_count(tokenizer: Tokenizer, rows: list[list[int]]) -> int:
+    """Count decoded UTF-8 bytes across every target-token row."""
+    return sum(len(tokenizer.decode(row).encode("utf-8")) for row in rows)
+
+
+def pad_or_trim(ids: list[int], length: int, pad_id: int) -> list[int]:
+    if len(ids) >= length:
+        return ids[:length]
+    return ids + [pad_id] * (length - len(ids))
+
+
+def encode_sft_example(
+    tokenizer: Tokenizer,
+    system_prompt: str,
+    user: str,
+    assistant: str,
+    sequence_length: int,
+) -> tuple[list[int], list[int]]:
+    pad_id = tokenizer.token_to_id(PAD) or 0
+    prefix_ids = encode_ids(tokenizer, format_prompt(system_prompt, user))
+    full_ids = encode_ids(tokenizer, format_chat(system_prompt, user, assistant))
+    model_tokens = sequence_length + 1
+    if len(full_ids) > model_tokens:
+        overflow = len(full_ids) - model_tokens
+        prefix_ids = prefix_ids[overflow:] if overflow < len(prefix_ids) else []
+        full_ids = full_ids[-model_tokens:]
+    input_ids = pad_or_trim(full_ids[:-1], sequence_length, pad_id)
+    target_ids = full_ids[1:]
+    first_assistant_target = max(0, min(len(prefix_ids) - 1, len(target_ids)))
+    labels = [IGNORE_INDEX] * first_assistant_target + target_ids[first_assistant_target:]
+    labels = pad_or_trim(labels, sequence_length, IGNORE_INDEX)
+    return input_ids, labels
+
+
+def window_tokens(ids: list[int], sequence_length: int) -> list[list[int]]:
+    if len(ids) < sequence_length + 1:
+        pad = [ids[-1] if ids else 0] * (sequence_length + 1 - len(ids))
+        ids = ids + pad
+    windows: list[list[int]] = []
+    stride = max(1, sequence_length // 2)
+    for start in range(0, len(ids) - sequence_length, stride):
+        windows.append(ids[start:start + sequence_length + 1])
+        if len(windows) >= 256:
+            break
+    if not windows:
+        windows.append(ids[:sequence_length + 1])
+    return windows
+
+
+def parse_numeric_answer(text: str) -> float | None:
+    matches = re.findall(r"-?\d+(?:\.\d+)?", text.replace(",", ""))
+    if not matches:
+        return None
+    return float(matches[-1])
+
+
+def numeric_reward(prediction: str, expected: str) -> float:
+    predicted = parse_numeric_answer(prediction)
+    target = parse_numeric_answer(expected)
+    if predicted is None or target is None:
+        return 0.0
+    return 1.0 if predicted == target else 0.0
+
+
+def example_pairs(raw: list[Any]) -> list[dict[str, str]]:
+    pairs: list[dict[str, str]] = []
+    for item in raw:
+        if not isinstance(item, dict):
+            continue
+        user = str(item.get("input") or item.get("user") or "")
+        assistant = str(item.get("output") or item.get("assistant") or "")
+        if user and assistant:
+            pairs.append({"input": user, "output": assistant})
+    return pairs

package/training/foundation/src/evaluate.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+import time
+from pathlib import Path
+
+import torch
+from tokenizers import Tokenizer
+from torch.nn import functional as F
+
+from common import load_config, write_json
+from data import END, decoded_byte_count, encode_ids, example_pairs, format_prompt, smoke_corpus, window_tokens
+from model import bits_per_byte, generate, load_model
+
+
+def pick_device() -> torch.device:
+    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def evaluate_bpb(model, tokenizer: Tokenizer, corpus: str, sequence_length: int, device: torch.device) -> dict[str, float | int]:
+    windows = window_tokens(tokenizer.encode(corpus).ids, sequence_length)
+    tensor = torch.tensor(windows[:8], dtype=torch.long, device=device)
+    inputs = tensor[:, :-1]
+    labels = tensor[:, 1:]
+    model.eval()
+    with torch.no_grad():
+        logits = model(inputs)
+        nll = F.cross_entropy(logits.reshape(-1, logits.size(-1)), labels.reshape(-1), reduction="sum")
+    byte_count = max(decoded_byte_count(tokenizer, labels.detach().cpu().tolist()), 1)
+    return {
+        "nll": float(nll.item()),
+        "bpb": bits_per_byte(nll, byte_count),
+        "tokens": int(inputs.numel()),
+        "bytes": int(byte_count),
+    }
+
+
+def evaluate_chat(
+    model,
+    tokenizer: Tokenizer,
+    system_prompt: str,
+    examples: list[dict[str, str]],
+    device: torch.device,
+) -> dict[str, float | int | list[dict[str, str | bool]]]:
+    rows: list[dict[str, str | bool]] = []
+    correct = 0
+    model.eval()
+    for example in examples:
+        prompt_ids = encode_ids(tokenizer, format_prompt(system_prompt, example["input"])) or [0]
+        tokens = torch.tensor([prompt_ids], dtype=torch.long, device=device)
+        generated = generate(
+            model,
+            tokens,
+            max_new_tokens=32,
+            stop_token_id=tokenizer.token_to_id(END),
+        )[0].tolist()[len(prompt_ids):]
+        text = tokenizer.decode(generated).split(END)[0].strip()
+        matched = text == example["output"].strip()
+        correct += int(matched)
+        rows.append({"input": example["input"], "expected": example["output"], "actual": text, "match": matched})
+    return {
+        "examples": len(examples),
+        "correct": correct,
+        "accuracy": (correct / len(examples)) if examples else 0.0,
+        "predictions": rows,
+    }
+
+
+def evaluate_inference(model, tokenizer: Tokenizer, prompt: str, device: torch.device) -> dict[str, float | int | str]:
+    prompt_ids = encode_ids(tokenizer, prompt) or [0]
+    tokens = torch.tensor([prompt_ids[: min(16, len(prompt_ids))]], dtype=torch.long, device=device)
+    if device.type == "cuda":
+        torch.cuda.reset_peak_memory_stats(device)
+        torch.cuda.synchronize()
+    started = time.perf_counter()
+    generated = generate(model, tokens, max_new_tokens=32)
+    if device.type == "cuda":
+        torch.cuda.synchronize()
+    elapsed = max(time.perf_counter() - started, 1e-6)
+    new_tokens = max(int(generated.size(1) - tokens.size(1)), 0)
+    vram_mb = float(torch.cuda.max_memory_allocated(device) / (1024 * 1024)) if device.type == "cuda" else 0.0
+    return {
+        "tokens": new_tokens,
+        "seconds": elapsed,
+        "tokens_per_second": new_tokens / elapsed,
+        "vram_mb": vram_mb,
+        "device": str(device),
+    }
+
+
+def main(argv: list[str] | None = None) -> None:
+    config = load_config(argv)
+    evaluator = str(config.get("evaluator") or "bpb")
+    device = pick_device()
+    tokenizer = Tokenizer.from_file(str(Path(config["tokenizer_dir"]) / "tokenizer.json"))
+    model = load_model(config["model_dir"], map_location=device)
+    examples = example_pairs(config.get("examples") or [])
+    sequence_length = int(config.get("sequence_length") or model.config["sequence_length"])
+    output_dir = Path(config["output_dir"])
+    metrics: dict = {"ok": True, "evaluator": evaluator, "device": str(device)}
+    if evaluator == "bpb":
+        corpus = smoke_corpus(examples, int(config.get("max_chars") or 20_000), extra=str(config.get("system_prompt") or ""))
+        metrics.update(evaluate_bpb(model, tokenizer, corpus, sequence_length, device))
+    elif evaluator == "chat":
+        metrics.update(evaluate_chat(model, tokenizer, str(config.get("system_prompt") or ""), examples, device))
+    elif evaluator == "inference":
+        prompt = examples[0]["input"] if examples else "hello"
+        metrics.update(evaluate_inference(model, tokenizer, prompt, device))
+    else:
+        raise SystemExit(f"Unknown foundation evaluator: {evaluator}")
+    write_json(output_dir / "metrics.json", metrics)
+    write_json(output_dir / "report.json", metrics)
+
+
+if __name__ == "__main__":
+    main()

package/training/foundation/src/finetune.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+from tokenizers import Tokenizer
+from torch.nn import functional as F
+
+from common import load_config, require_cuda, write_json
+from data import IGNORE_INDEX, encode_sft_example, example_pairs
+from model import load_model, param_count, save_model
+
+
+def sft_tensors(
+    tokenizer: Tokenizer,
+    system_prompt: str,
+    examples: list[dict[str, str]],
+    sequence_length: int,
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    inputs: list[list[int]] = []
+    labels: list[list[int]] = []
+    for example in examples:
+        input_ids, label_ids = encode_sft_example(
+            tokenizer,
+            system_prompt,
+            example["input"],
+            example["output"],
+            sequence_length,
+        )
+        inputs.append(input_ids)
+        labels.append(label_ids)
+    return (
+        torch.tensor(inputs, dtype=torch.long, device=device),
+        torch.tensor(labels, dtype=torch.long, device=device),
+    )
+
+
+def main(argv: list[str] | None = None) -> None:
+    config = load_config(argv)
+    require_cuda("finetune")
+    device = torch.device("cuda")
+    tokenizer = Tokenizer.from_file(str(Path(config["tokenizer_dir"]) / "tokenizer.json"))
+    model = load_model(config["model_dir"], map_location=device)
+    examples = example_pairs(config.get("examples") or [])
+    if len(examples) < 1:
+        raise SystemExit("Foundation finetune needs at least one chat example.")
+    sequence_length = int(config.get("sequence_length") or model.config["sequence_length"])
+    inputs, labels = sft_tensors(
+        tokenizer,
+        str(config.get("system_prompt") or ""),
+        examples,
+        sequence_length,
+        device,
+    )
+    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
+    steps = max(1, int(config.get("steps") or 2))
+    batch_size = max(1, int(config.get("batch_size") or 2))
+    last_loss = 0.0
+    model.train()
+    for step in range(steps):
+        start = (step * batch_size) % inputs.size(0)
+        end = start + batch_size
+        if end <= inputs.size(0):
+            batch_inputs = inputs[start:end]
+            batch_labels = labels[start:end]
+        else:
+            indices = [(start + offset) % inputs.size(0) for offset in range(batch_size)]
+            batch_inputs = inputs[indices]
+            batch_labels = labels[indices]
+        logits = model(batch_inputs)
+        loss = F.cross_entropy(
+            logits.reshape(-1, logits.size(-1)),
+            batch_labels.reshape(-1),
+            ignore_index=IGNORE_INDEX,
+        )
+        optimizer.zero_grad(set_to_none=True)
+        loss.backward()
+        optimizer.step()
+        last_loss = float(loss.item())
+
+    output_dir = Path(config["output_dir"])
+    save_model(model, output_dir)
+    write_json(output_dir / "metrics.json", {
+        "ok": True,
+        "steps": steps,
+        "loss": last_loss,
+        "examples": inputs.size(0),
+        "parameters": param_count(model),
+        "ignore_index": IGNORE_INDEX,
+        "device": str(device),
+    })
+
+
+if __name__ == "__main__":
+    main()

package/training/foundation/src/model.py

@@ -0,0 +1,157 @@
+from __future__ import annotations
+
+import json
+import math
+from pathlib import Path
+from typing import Any
+
+import torch
+from torch import nn
+from safetensors.torch import load_file, save_file
+
+from common import ensure_private_directory, make_private_file, write_json
+
+
+def derived_width(depth: int) -> int:
+    return 32 * int(depth)
+
+
+def derived_heads(depth: int) -> int:
+    width = derived_width(depth)
+    for heads in (8, 4, 2, 1):
+        if width % heads == 0:
+            return heads
+    return 1
+
+
+def model_config_from_depth(depth: int, vocab_size: int, sequence_length: int) -> dict[str, int]:
+    return {
+        "depth": int(depth),
+        "width": derived_width(depth),
+        "heads": derived_heads(depth),
+        "vocab_size": int(vocab_size),
+        "sequence_length": int(sequence_length),
+    }
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, width: int, heads: int) -> None:
+        super().__init__()
+        if width % heads != 0:
+            raise ValueError("width must be divisible by heads")
+        self.heads = heads
+        self.head_dim = width // heads
+        self.qkv = nn.Linear(width, 3 * width, bias=False)
+        self.proj = nn.Linear(width, width, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        batch, steps, width = x.shape
+        qkv = self.qkv(x).view(batch, steps, 3, self.heads, self.head_dim)
+        query, key, value = qkv.unbind(dim=2)
+        query = query.transpose(1, 2)
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
+        context = torch.nn.functional.scaled_dot_product_attention(
+            query, key, value, is_causal=True,
+        )
+        context = context.transpose(1, 2).contiguous().view(batch, steps, width)
+        return self.proj(context)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, width: int, heads: int) -> None:
+        super().__init__()
+        self.ln1 = nn.LayerNorm(width)
+        self.attn = CausalSelfAttention(width, heads)
+        self.ln2 = nn.LayerNorm(width)
+        self.mlp = nn.Sequential(
+            nn.Linear(width, 4 * width),
+            nn.GELU(),
+            nn.Linear(4 * width, width),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.attn(self.ln1(x))
+        x = x + self.mlp(self.ln2(x))
+        return x
+
+
+class FoundationGPT(nn.Module):
+    def __init__(self, config: dict[str, int]) -> None:
+        super().__init__()
+        self.config = config
+        width = config["width"]
+        self.tok_emb = nn.Embedding(config["vocab_size"], width)
+        self.pos_emb = nn.Embedding(config["sequence_length"], width)
+        self.blocks = nn.ModuleList(
+            [TransformerBlock(width, config["heads"]) for _ in range(config["depth"])],
+        )
+        self.ln_f = nn.LayerNorm(width)
+        self.lm_head = nn.Linear(width, config["vocab_size"], bias=False)
+        self.lm_head.weight = self.tok_emb.weight
+
+    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
+        _batch, steps = input_ids.shape
+        if steps > self.config["sequence_length"]:
+            raise ValueError("sequence exceeds model sequence_length")
+        positions = torch.arange(steps, device=input_ids.device)
+        x = self.tok_emb(input_ids) + self.pos_emb(positions)
+        for block in self.blocks:
+            x = block(x)
+        return self.lm_head(self.ln_f(x))
+
+
+@torch.no_grad()
+def generate(
+    model: FoundationGPT,
+    input_ids: torch.Tensor,
+    max_new_tokens: int,
+    stop_token_id: int | None = None,
+) -> torch.Tensor:
+    model.eval()
+    tokens = input_ids
+    limit = model.config["sequence_length"]
+    finished = torch.zeros(tokens.size(0), dtype=torch.bool, device=tokens.device)
+    for _ in range(max_new_tokens):
+        logits = model(tokens[:, -limit:])
+        next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True)
+        if stop_token_id is not None:
+            next_token = torch.where(
+                finished.unsqueeze(1),
+                torch.full_like(next_token, stop_token_id),
+                next_token,
+            )
+        tokens = torch.cat([tokens, next_token], dim=1)
+        if stop_token_id is not None:
+            finished |= next_token.squeeze(1).eq(stop_token_id)
+            if bool(finished.all()):
+                break
+    return tokens
+
+
+def save_model(model: FoundationGPT, directory: str | Path) -> None:
+    destination = ensure_private_directory(directory)
+    tensors = {key: value.detach().cpu().contiguous() for key, value in model.state_dict().items()}
+    model_path = destination / "model.safetensors"
+    save_file(tensors, str(model_path))
+    make_private_file(model_path)
+    write_json(destination / "config.json", model.config)
+
+
+def load_model(directory: str | Path, map_location: str | torch.device = "cpu") -> FoundationGPT:
+    destination = Path(directory)
+    config: dict[str, Any] = json.loads((destination / "config.json").read_text(encoding="utf-8"))
+    model = FoundationGPT({key: int(value) for key, value in config.items()})
+    state = load_file(str(destination / "model.safetensors"))
+    model.load_state_dict(state)
+    return model.to(map_location)
+
+
+def param_count(model: nn.Module) -> int:
+    return sum(parameter.numel() for parameter in model.parameters())
+
+
+def bits_per_byte(nll_sum: torch.Tensor, byte_count: int) -> float:
+    if byte_count <= 0:
+        return 0.0
+    return float(nll_sum.item() / math.log(2) / byte_count)