@tuned-tensor/cli 0.13.1-beta.5 → 0.13.1-beta.7

package/docs/local-runtime/architecture.md

@@ -51,19 +51,37 @@ TT Local currently ships two such methods:
   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
+The foundation engine is a readable single-GPU trainer inspired by nanochat,
+not a port of nanochat's distributed or compute-optimal training stack. It can
+stream configured text/JSONL shards into a deterministic token cache, use a
+distinct held-out BPB corpus, and atomically recover model, optimizer,
+scheduler, RNG, counters, and token position. Without a configured corpus it
+retains the bounded prompt/example smoke path. Chat evaluation reuses the spec
+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.
+example-derived unless `validation_path` is configured. Smoke 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.
 
+## Spec-driven instruction contract
+
+Adapter and foundation specs extend one engine-neutral behavior schema:
+`name`, `description`, `system_prompt`, `guidelines`, `constraints`, and
+`examples`. Node compiles the three instruction fields into one canonical
+system message. The adapter dataset compiler and every foundation stage use
+that same compiler, so tokenizer input, pretraining smoke data, chat SFT, RL,
+and evaluation all receive guidelines and constraints consistently.
+
+Engine-specific fields control execution only. Adapter specs add `base_model`,
+LoRA hyperparameters, and optional prebuilt datasets. Foundation specs add
+`engine: "foundation"` and the `foundation` training configuration. The
+behavioral instruction remains the shared source of truth.
+
 ## Python boundary
 
 Node owns input validation, state, process lifecycle, manifests, and reports.

package/docs/local-runtime/foundation-long-runs.md

@@ -0,0 +1,108 @@
+# Long foundation runs
+
+The foundation engine can run recoverable, single-GPU pretraining from `.txt`
+files or JSONL shards whose rows contain a string `text` field. Keep training
+and validation data in different files or directories.
+
+## Spec
+
+The long-run fields live under `foundation` in `tunedtensor.json`:
+
+```json
+{
+  "corpus_path": "./data/train",
+  "validation_path": "./data/validation.jsonl",
+  "checkpoint_backup_dir": "/mnt/checkpoint-backup/my-run",
+  "depth": 20,
+  "vocab_size": 32000,
+  "tokenizer_max_chars": 100000000,
+  "max_chars": 50000000000,
+  "pretrain_steps": 200000,
+  "sequence_length": 1024,
+  "batch_size": 2,
+  "gradient_accumulation_steps": 16,
+  "learning_rate": 0.0003,
+  "weight_decay": 0.1,
+  "warmup_steps": 2000,
+  "min_lr_ratio": 0.1,
+  "gradient_clip": 1.0,
+  "bf16": true,
+  "checkpoint_interval_steps": 500,
+  "checkpoint_interval_seconds": 1800,
+  "keep_checkpoints": 3,
+  "log_interval_steps": 10,
+  "seed": 1337,
+  "finetune_steps": 100,
+  "rl_steps": 0,
+  "nproc_per_node": 1
+}
+```
+
+Relative paths are resolved beside the spec. `max_chars` is a deterministic
+upper bound across the sorted corpus shards. Omit `corpus_path` only for the
+small prompt/example smoke corpus.
+
+Run the preflight, then use one stable, absolute run directory from the first
+launch onward:
+
+```bash
+tt doctor tunedtensor.json
+tt pipeline run \
+  --spec /absolute/project/tunedtensor.json \
+  --resume /absolute/project/.tuned-tensor/foundation-runs/experiment-1
+```
+
+`--resume` also starts a missing directory. On restart, verified completed
+stages are skipped. An interrupted pretrain stage restores the newest complete
+checkpoint with matching configuration. Changing model, optimizer, tokenizer,
+data manifest, or schedule settings requires a new run directory.
+
+Recovery data is stored under `pretrain/recovery/`:
+
+- `pretrain-tokens.bin` and its source/tokenizer manifest;
+- `training-metrics.jsonl` heartbeats;
+- `checkpoints/step-*` atomic rolling checkpoints.
+
+Each checkpoint contains model weights, AdamW and scheduler state, CPU and CUDA
+RNG state, completed optimizer step, samples consumed, token count, and loss.
+SIGINT or SIGTERM requests a step-boundary checkpoint. The CLI allows up to two
+minutes for that checkpoint before forcibly stopping the process. SIGKILL or a
+power loss falls back to the most recent periodic checkpoint.
+
+## User service
+
+For unattended operation, run the resumable command under a user systemd
+service. Replace every path with an absolute path and use the installed `tt`
+path reported by `command -v tt`:
+
+```ini
+[Unit]
+Description=Tuned Tensor foundation pretraining
+After=network-online.target
+
+[Service]
+Type=simple
+WorkingDirectory=/absolute/project
+ExecStart=/absolute/path/to/tt pipeline run --spec /absolute/project/tunedtensor.json --resume /absolute/project/.tuned-tensor/foundation-runs/experiment-1
+Restart=on-failure
+RestartSec=30
+TimeoutStopSec=130
+KillSignal=SIGTERM
+
+[Install]
+WantedBy=default.target
+```
+
+Save it as `~/.config/systemd/user/tt-foundation.service`, then run:
+
+```bash
+systemctl --user daemon-reload
+systemctl --user enable --now tt-foundation.service
+journalctl --user -fu tt-foundation.service
+```
+
+Enable user lingering if the job must survive logout. Use a UPS for power-loss
+protection, and put `checkpoint_backup_dir` on a different physical disk or
+remote mounted filesystem. Before a multi-day run, complete a full-load soak,
+stop the service during pretraining, and verify the journal reports a resumed
+checkpoint after restart.

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@tuned-tensor/cli",
-  "version": "0.13.1-beta.5",
+  "version": "0.13.1-beta.7",
   "description": "Local TT agent and CUDA fine-tuning CLI for Tuned Tensor workflows",
   "repository": {
     "type": "git",

package/training/foundation/src/checkpoint.py

@@ -0,0 +1,235 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+import hashlib
+import json
+import os
+from pathlib import Path
+import shutil
+import tempfile
+from typing import Any
+
+import torch
+from safetensors.torch import load_file, save_file
+
+from common import ensure_private_directory, write_json
+
+
+def _sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as source:
+        while chunk := source.read(1_048_576):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+@dataclass(frozen=True)
+class ResumeState:
+    step: int
+    samples_consumed: int
+    tokens_seen: int
+    last_loss: float
+    path: Path
+
+
+class CheckpointManager:
+    def __init__(
+        self,
+        output_dir: str | Path,
+        config_fingerprint: str,
+        *,
+        keep: int = 3,
+        backup_dir: str | Path | None = None,
+    ) -> None:
+        self.root = ensure_private_directory(Path(output_dir) / "checkpoints").resolve()
+        self.config_fingerprint = config_fingerprint
+        self.keep = max(1, int(keep))
+        self.backup_dir = Path(backup_dir).expanduser().resolve() if backup_dir else None
+        if self.backup_dir is not None and (
+            self.backup_dir == self.root or self.root in self.backup_dir.parents
+        ):
+            raise ValueError("checkpoint_backup_dir must be outside the primary checkpoints directory.")
+
+    @staticmethod
+    def _name(step: int) -> str:
+        return f"step-{step:09d}"
+
+    def _candidates(self, root: Path | None = None) -> list[Path]:
+        candidate_root = root or self.root
+        return sorted(
+            (path for path in candidate_root.glob("step-*") if path.is_dir()),
+            key=lambda path: path.name,
+            reverse=True,
+        )
+
+    def _resume_candidates(self) -> list[Path]:
+        ranked = [(path, 1) for path in self._candidates()]
+        if self.backup_dir is not None and self.backup_dir.is_dir():
+            ranked.extend((path, 0) for path in self._candidates(self.backup_dir))
+        return [
+            path
+            for path, _priority in sorted(
+                ranked,
+                key=lambda candidate: (candidate[0].name, candidate[1]),
+                reverse=True,
+            )
+        ]
+
+    def latest(self) -> Path | None:
+        mismatched = False
+        corrupt = False
+        for path in self._resume_candidates():
+            try:
+                metadata = json.loads((path / "metadata.json").read_text(encoding="utf-8"))
+            except (OSError, json.JSONDecodeError):
+                corrupt = True
+                continue
+            if not isinstance(metadata, dict) or metadata.get("complete") is not True:
+                corrupt = True
+                continue
+            if metadata.get("config_fingerprint") != self.config_fingerprint:
+                mismatched = True
+                continue
+            model_path = path / "model.safetensors"
+            state_path = path / "training-state.pt"
+            if not model_path.is_file() or not state_path.is_file():
+                corrupt = True
+                continue
+            files = metadata.get("files")
+            if not isinstance(files, dict):
+                corrupt = True
+                continue
+            valid = True
+            for name, candidate in [
+                ("model.safetensors", model_path),
+                ("training-state.pt", state_path),
+            ]:
+                expected = files.get(name)
+                if (
+                    not isinstance(expected, dict)
+                    or candidate.stat().st_size != expected.get("size")
+                    or _sha256(candidate) != expected.get("sha256")
+                ):
+                    valid = False
+                    break
+            if valid:
+                return path
+            corrupt = True
+        if mismatched:
+            raise ValueError("Existing foundation checkpoints do not match the current training configuration.")
+        if corrupt:
+            raise ValueError("No intact foundation checkpoint remains for the current training configuration.")
+        return None
+
+    def load(
+        self,
+        model: torch.nn.Module,
+        optimizer: torch.optim.Optimizer,
+        scheduler: torch.optim.lr_scheduler.LRScheduler,
+        device: torch.device,
+    ) -> ResumeState | None:
+        path = self.latest()
+        if path is None:
+            return None
+        model.load_state_dict(load_file(str(path / "model.safetensors"), device=str(device)))
+        state = torch.load(path / "training-state.pt", map_location=device, weights_only=True)
+        optimizer.load_state_dict(state["optimizer"])
+        scheduler.load_state_dict(state["scheduler"])
+        torch.set_rng_state(state["cpu_rng"].cpu())
+        if device.type == "cuda" and state.get("cuda_rng"):
+            torch.cuda.set_rng_state_all([value.cpu() for value in state["cuda_rng"]])
+        return ResumeState(
+            step=int(state["step"]),
+            samples_consumed=int(state["samples_consumed"]),
+            tokens_seen=int(state["tokens_seen"]),
+            last_loss=float(state.get("last_loss", 0.0)),
+            path=path,
+        )
+
+    def save(
+        self,
+        model: torch.nn.Module,
+        optimizer: torch.optim.Optimizer,
+        scheduler: torch.optim.lr_scheduler.LRScheduler,
+        *,
+        step: int,
+        samples_consumed: int,
+        tokens_seen: int,
+        last_loss: float,
+        reason: str,
+    ) -> Path:
+        name = self._name(step)
+        final = self.root / name
+        if final.exists() and self.latest() == final:
+            self._backup(final)
+            return final
+        temporary = Path(tempfile.mkdtemp(prefix=f".{name}.", dir=self.root))
+        temporary.chmod(0o700)
+        try:
+            tensors = {key: value.detach().cpu().contiguous() for key, value in model.state_dict().items()}
+            model_path = temporary / "model.safetensors"
+            save_file(tensors, str(model_path))
+            model_path.chmod(0o600)
+            state_path = temporary / "training-state.pt"
+            torch.save({
+                "step": int(step),
+                "samples_consumed": int(samples_consumed),
+                "tokens_seen": int(tokens_seen),
+                "last_loss": float(last_loss),
+                "optimizer": optimizer.state_dict(),
+                "scheduler": scheduler.state_dict(),
+                "cpu_rng": torch.get_rng_state(),
+                "cuda_rng": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else [],
+            }, state_path)
+            state_path.chmod(0o600)
+            write_json(temporary / "metadata.json", {
+                "complete": True,
+                "config_fingerprint": self.config_fingerprint,
+                "step": int(step),
+                "samples_consumed": int(samples_consumed),
+                "tokens_seen": int(tokens_seen),
+                "reason": reason,
+                "files": {
+                    "model.safetensors": {
+                        "size": model_path.stat().st_size,
+                        "sha256": _sha256(model_path),
+                    },
+                    "training-state.pt": {
+                        "size": state_path.stat().st_size,
+                        "sha256": _sha256(state_path),
+                    },
+                },
+            })
+            if final.exists():
+                shutil.rmtree(final)
+            os.replace(temporary, final)
+            self._prune()
+            self._backup(final)
+            return final
+        finally:
+            if temporary.exists():
+                shutil.rmtree(temporary, ignore_errors=True)
+
+    def _prune(self) -> None:
+        for stale in self._candidates()[self.keep:]:
+            shutil.rmtree(stale)
+
+    def _backup(self, checkpoint: Path) -> None:
+        if self.backup_dir is None:
+            return
+        ensure_private_directory(self.backup_dir)
+        destination = self.backup_dir / checkpoint.name
+        temporary = self.backup_dir / f".{checkpoint.name}.{os.getpid()}.tmp"
+        if temporary.exists():
+            shutil.rmtree(temporary)
+        shutil.copytree(checkpoint, temporary)
+        if destination.exists():
+            shutil.rmtree(destination)
+        os.replace(temporary, destination)
+        retained = sorted(
+            (path for path in self.backup_dir.glob("step-*") if path.is_dir()),
+            key=lambda path: path.name,
+            reverse=True,
+        )
+        for stale in retained[self.keep:]:
+            shutil.rmtree(stale)

package/training/foundation/src/common.py

@@ -2,7 +2,9 @@ from __future__ import annotations
 
 import argparse
 import json
+import os
 from pathlib import Path
+import tempfile
 from typing import Any
 
 
@@ -29,9 +31,35 @@ def load_config(argv: list[str] | None = None) -> dict[str, Any]:
 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:
+    descriptor, temporary_name = tempfile.mkstemp(
+        dir=destination.parent,
+        prefix=f".{destination.name}.",
+        suffix=".tmp",
+    )
+    temporary = Path(temporary_name)
+    try:
+        with os.fdopen(descriptor, "w", encoding="utf-8") as output:
+            os.fchmod(output.fileno(), 0o600)
+            output.write(json.dumps(value, indent=2) + "\n")
+            output.flush()
+            os.fsync(output.fileno())
+        os.replace(temporary, destination)
         make_private_file(destination)
-        output.write(json.dumps(value, indent=2) + "\n")
+    finally:
+        temporary.unlink(missing_ok=True)
+
+
+def append_jsonl(path: str | Path, value: Any) -> None:
+    destination = Path(path)
+    ensure_private_directory(destination.parent)
+    descriptor = os.open(destination, os.O_WRONLY | os.O_CREAT | os.O_APPEND, 0o600)
+    try:
+        payload = (json.dumps(value, separators=(",", ":")) + "\n").encode("utf-8")
+        os.write(descriptor, payload)
+        os.fsync(descriptor)
+    finally:
+        os.close(descriptor)
+    make_private_file(destination)
 
 
 def require_cuda(action: str) -> None:

package/training/foundation/src/data.py

@@ -1,8 +1,11 @@
 from __future__ import annotations
 
 import hashlib
+import json
+import os
+from pathlib import Path
 import re
-from typing import Any
+from typing import Any, Iterable, Iterator
 
 from tokenizers import Tokenizer
 from tokenizers import models as tokenizer_models
@@ -18,7 +21,7 @@ SPECIAL_TOKENS = [UNK, PAD, SYSTEM, USER, ASSISTANT, END]
 IGNORE_INDEX = -100
 
 
-def train_tokenizer(corpus: str, vocab_size: int) -> Tokenizer:
+def train_tokenizer(corpus: str | Iterable[str], vocab_size: int) -> Tokenizer:
     tokenizer = Tokenizer(tokenizer_models.BPE(unk_token=UNK))
     tokenizer.pre_tokenizer = pre_tokenizers.Whitespace()
     trainer = trainers.BpeTrainer(
@@ -26,10 +29,100 @@ def train_tokenizer(corpus: str, vocab_size: int) -> Tokenizer:
         special_tokens=SPECIAL_TOKENS,
         min_frequency=1,
     )
-    tokenizer.train_from_iterator([corpus], trainer=trainer)
+    tokenizer.train_from_iterator([corpus] if isinstance(corpus, str) else corpus, trainer=trainer)
     return tokenizer
 
 
+def corpus_files(path: str | Path) -> list[Path]:
+    unresolved = Path(path).expanduser()
+    if unresolved.is_symlink():
+        raise ValueError(f"Foundation corpus must not be a symbolic link: {unresolved}")
+    source = unresolved.resolve()
+    if source.is_file():
+        files = [source]
+    elif source.is_dir():
+        files = sorted(
+            candidate for candidate in source.rglob("*")
+            if candidate.is_file() and candidate.suffix.lower() in {".txt", ".jsonl"}
+        )
+    else:
+        raise ValueError(f"Foundation corpus path does not exist: {source}")
+    if not files:
+        raise ValueError(f"Foundation corpus has no .txt or .jsonl files: {source}")
+    for candidate in files:
+        if candidate.is_symlink():
+            raise ValueError(f"Foundation corpus must not contain symbolic links: {candidate}")
+    return files
+
+
+def corpus_manifest(path: str | Path) -> list[dict[str, int | str]]:
+    manifest: list[dict[str, int | str]] = []
+    for candidate in corpus_files(path):
+        info = candidate.stat()
+        manifest.append({
+            "path": str(candidate),
+            "size": int(info.st_size),
+            "mtime_ns": int(info.st_mtime_ns),
+        })
+    return manifest
+
+
+def iter_corpus_documents(path: str | Path, max_chars: int | None = None) -> Iterator[str]:
+    remaining = max_chars if max_chars and max_chars > 0 else None
+    for candidate in corpus_files(path):
+        if remaining is not None and remaining <= 0:
+            return
+        if candidate.suffix.lower() == ".jsonl":
+            with candidate.open("r", encoding="utf-8") as source:
+                for line_number, line in enumerate(source, start=1):
+                    if remaining is not None and remaining <= 0:
+                        return
+                    if not line.strip():
+                        continue
+                    try:
+                        row = json.loads(line)
+                    except json.JSONDecodeError as error:
+                        raise ValueError(f"Invalid JSONL at {candidate}:{line_number}: {error.msg}") from error
+                    if not isinstance(row, dict) or not isinstance(row.get("text"), str):
+                        raise ValueError(f'Foundation JSONL rows require a string "text" field: {candidate}:{line_number}')
+                    text = row["text"]
+                    if remaining is not None:
+                        text = text[:remaining]
+                        remaining -= len(text)
+                    if text:
+                        yield text
+        else:
+            with candidate.open("r", encoding="utf-8") as source:
+                while True:
+                    requested = min(1_048_576, remaining) if remaining is not None else 1_048_576
+                    if requested <= 0:
+                        return
+                    text = source.read(requested)
+                    if not text:
+                        break
+                    if remaining is not None:
+                        remaining -= len(text)
+                    yield text
+
+
+def instruction_corpus_documents(system_prompt: str, documents: Iterable[str]) -> Iterator[str]:
+    instruction = system_prompt.strip()
+    if instruction:
+        yield instruction
+    yield from documents
+
+
+def corpus_documents_hash(documents: Iterable[str]) -> tuple[str, int]:
+    digest = hashlib.sha256()
+    chars = 0
+    for document in documents:
+        encoded = document.encode("utf-8")
+        digest.update(len(encoded).to_bytes(8, "little"))
+        digest.update(encoded)
+        chars += len(document)
+    return digest.hexdigest(), chars
+
+
 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"
@@ -55,6 +148,52 @@ def encode_ids(tokenizer: Tokenizer, text: str) -> list[int]:
     return list(tokenizer.encode(text).ids)
 
 
+def encode_prompt_ids(
+    tokenizer: Tokenizer,
+    system_prompt: str,
+    user: str,
+    max_tokens: int,
+) -> list[int]:
+    """Keep the complete spec instruction and trim only user content."""
+    full_ids = encode_ids(tokenizer, format_prompt(system_prompt, user))
+    if len(full_ids) <= max_tokens:
+        return full_ids
+    system = system_prompt.strip() or "You are a helpful assistant."
+    prefix_ids = encode_ids(tokenizer, f"{SYSTEM}{system}{END}{USER}")
+    suffix_ids = encode_ids(tokenizer, f"{END}{ASSISTANT}")
+    if len(prefix_ids) + len(suffix_ids) > max_tokens:
+        raise ValueError(
+            "Compiled spec instruction exceeds the foundation model context; "
+            "increase foundation.sequence_length or shorten the spec instruction."
+        )
+    user_ids = encode_ids(tokenizer, user)
+    user_budget = max_tokens - len(prefix_ids) - len(suffix_ids)
+    retained_user_ids = user_ids[-user_budget:] if user_budget > 0 else []
+    return prefix_ids + retained_user_ids + suffix_ids
+
+
+def encode_generation_prompt(
+    tokenizer: Tokenizer,
+    system_prompt: str,
+    user: str,
+    context_length: int,
+    requested_completion_tokens: int,
+) -> tuple[list[int], int]:
+    """Reserve completion space without rejecting an instruction that can fit."""
+    completion_tokens = min(requested_completion_tokens, max(1, context_length // 2))
+    try:
+        prompt_ids = encode_prompt_ids(
+            tokenizer,
+            system_prompt,
+            user,
+            context_length - completion_tokens,
+        )
+    except ValueError:
+        prompt_ids = encode_prompt_ids(tokenizer, system_prompt, user, context_length - 1)
+        completion_tokens = context_length - len(prompt_ids)
+    return prompt_ids, completion_tokens
+
+
 def decode_ids(tokenizer: Tokenizer, ids: list[int]) -> str:
     return tokenizer.decode(ids, skip_special_tokens=False)
 
@@ -82,9 +221,18 @@ def encode_sft_example(
     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:]
+        assistant_ids = encode_ids(tokenizer, f"{assistant}{END}")
+        reserved_assistant_tokens = min(2, len(assistant_ids))
+        prefix_ids = encode_prompt_ids(
+            tokenizer,
+            system_prompt,
+            user,
+            model_tokens - reserved_assistant_tokens,
+        )
+        assistant_budget = model_tokens - len(prefix_ids)
+        if len(assistant_ids) > assistant_budget:
+            assistant_ids = assistant_ids[:assistant_budget - 1] + assistant_ids[-1:]
+        full_ids = prefix_ids + assistant_ids
     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)))