@tuned-tensor/cli 0.20.0 → 0.22.0

package/training/adapter/src/prefetch_dataset.py

@@ -0,0 +1,116 @@
+"""Prefetch the dataset files a behavior spec trains on from the Hugging Face Hub.
+
+Mirrors prefetch.py for models: the files land in the shared Hugging Face cache
+(``<HF_HOME>/hub/datasets--org--name``), the revision is pinned to an immutable
+commit, and every requested file is checked against its content-addressed blob.
+Only the split files named in the spec are downloaded, never the whole repo.
+"""
+import argparse
+import json
+import re
+from pathlib import Path
+from typing import Any
+
+from hub_download import configure_download_environment, download_snapshot, verify_blob_digest
+from prefetch import configure_hugging_face_cache, write_json
+
+HUGGING_FACE_REPO_ID = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]{0,95}/[A-Za-z0-9][A-Za-z0-9._-]{0,95}")
+COMMIT_SHA = re.compile(r"[0-9a-fA-F]{40}")
+
+
+def validate_repo_path(value: Any) -> str:
+    if not isinstance(value, str) or not value or value.startswith(("/", "\\")):
+        raise ValueError(f"Dataset file must be a repository-relative path: {value!r}")
+    parts = re.split(r"[/\\]", value)
+    if any(part in ("", ".", "..") for part in parts) or re.match(r"^[A-Za-z]:", value):
+        raise ValueError(f"Dataset file must be a repository-relative path: {value!r}")
+    return "/".join(parts)
+
+
+def verify_dataset_files(snapshot: Path, files: dict[str, str]) -> dict[str, dict[str, Any]]:
+    """Check each split file is a complete cache entry inside the dataset repo.
+
+    Normal caches link snapshot files to content-addressed blobs, which are
+    checked against their digest. Caches without symlink support (Windows
+    without Developer Mode) hold plain copies; those are contained and hashed.
+    """
+    repository = snapshot.parent.parent.resolve()
+    verified: dict[str, dict[str, Any]] = {}
+    for split, relative in files.items():
+        path = snapshot / relative
+        if not path.is_file():
+            raise ValueError(f"Dataset file {relative} ({split}) is missing from snapshot {snapshot}")
+        target = path.resolve()
+        if repository not in target.parents:
+            raise ValueError(f"Dataset file {relative} escapes the Hugging Face dataset cache: {target}")
+        size = target.stat().st_size
+        if size == 0:
+            raise ValueError(f"Dataset file {relative} ({split}) is empty")
+        _, sha256 = verify_blob_digest(path)
+        verified[split] = {
+            "path": relative,
+            "local_path": str(path),
+            "size_bytes": size,
+            "sha256": sha256,
+        }
+    return verified
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Prefetch Hugging Face dataset files for a TT behavior spec.")
+    parser.add_argument("--input", required=True)
+    parser.add_argument("--output", required=True)
+    args = parser.parse_args()
+
+    payload = json.loads(Path(args.input).read_text(encoding="utf-8"))
+    repo = str(payload["repo"])
+    if not HUGGING_FACE_REPO_ID.fullmatch(repo):
+        raise ValueError(f"Dataset repo must be a Hugging Face id like org/name: {repo!r}")
+    revision = payload.get("revision")
+    if revision is not None and not COMMIT_SHA.fullmatch(str(revision)):
+        raise ValueError("Dataset revision must be a 40-character Hugging Face commit SHA")
+    raw_files = payload.get("files")
+    if not isinstance(raw_files, dict) or not raw_files:
+        raise ValueError("Dataset prefetch requires at least one split file")
+    files = {str(split): validate_repo_path(path) for split, path in raw_files.items()}
+
+    configure_hugging_face_cache(payload.get("model_cache"))
+    configure_download_environment()
+    from huggingface_hub import constants
+
+    local_only = bool(payload.get("local_files_only", False))
+    action = "Verifying cached" if local_only else "Prefetching"
+    print(f"{action} dataset {repo} in {constants.HF_HUB_CACHE}...", flush=True)
+    snapshot_path, commit = download_snapshot(
+        repo_id=repo,
+        repo_type="dataset",
+        revision=revision,
+        allow_patterns=sorted(set(files.values())),
+        local_files_only=local_only,
+        label="dataset_prefetch",
+    )
+    snapshot = Path(snapshot_path)
+    if not COMMIT_SHA.fullmatch(snapshot.name):
+        raise ValueError("Hugging Face dataset snapshot did not resolve to a 40-character immutable commit SHA")
+    if revision and snapshot.name.lower() != str(revision).lower():
+        raise ValueError(f"Hugging Face returned dataset revision {snapshot.name}, not requested revision {revision}")
+    print(f"Verifying dataset files under {snapshot}...", flush=True)
+    verified = verify_dataset_files(snapshot, files)
+
+    write_json(args.output, {
+        "ok": True,
+        "repo": repo,
+        "repo_type": "dataset",
+        "requested_revision": revision,
+        "snapshot_revision": snapshot.name.lower(),
+        "snapshot_path": str(snapshot),
+        "hf_home": str(constants.HF_HOME),
+        "hub_cache": str(constants.HF_HUB_CACHE),
+        "files": verified,
+        "size_bytes": sum(item["size_bytes"] for item in verified.values()),
+        "commit": commit,
+    })
+
+
+if __name__ == "__main__":
+    main()

package/training/adapter/src/serve.py

@@ -11,7 +11,11 @@ from pathlib import Path
 from tempfile import TemporaryDirectory
 
 from evaluate import configure_hugging_face_cache, resolve_adapter_path
-from model_contract import assert_certified_base_model_revision, assert_certified_model_config
+from model_contract import (
+    assert_certified_base_model_revision,
+    assert_certified_model_config,
+    is_certified_base_model,
+)
 
 class BearerAuthMiddleware:
     """Keep TT's auth boundary on ALL routes, including upstream admin APIs.
@@ -38,6 +42,27 @@ class BearerAuthMiddleware:
         await self.app(scope, receive, send)
 
 
+def native_chat_template(model_source: str) -> str:
+    source = Path(model_source)
+    template_file = source / "chat_template.jinja"
+    template = (template_file.read_text(encoding="utf-8") if template_file.is_file()
+                else json.loads((source / "tokenizer_config.json").read_text())["chat_template"])
+    if not isinstance(template, str):
+        raise ValueError("Serving requires an unambiguous string chat template.")
+    return template
+
+
+def pin_native_content_format(args: list[str], model_source: str) -> None:
+    if "--chat-template" not in args:
+        return
+    # The wrapper handles system content parts, which confuses upstream's AST
+    # detection for string-only native templates. Detect the ORIGINAL template
+    # using the detector from our pinned vLLM runtime, then pin that format.
+    from vllm.renderers.hf import _detect_content_format
+    content_format = _detect_content_format(native_chat_template(model_source), default="string")
+    args += ["--chat-template-content-format", content_format]
+
+
 def build_vllm_args(model_source: str, adapter_path: str | None, temp_dir: Path) -> list[str]:
     args = [
         "serve", model_source,
@@ -58,6 +83,8 @@ def build_vllm_args(model_source: str, adapter_path: str | None, temp_dir: Path)
     ]
     if os.environ.get("TT_API_KEY"):
         args += ["--middleware", "serve.BearerAuthMiddleware"]
+    if not is_certified_base_model(os.environ.get("TT_BASE_MODEL", "")):
+        args += ["--load-format", "safetensors"]
     base = os.environ["TT_BASE_MODEL"]
     # Upstream model-specific parsers, not a TT parser or model-size heuristic.
     parser = {
@@ -81,14 +108,9 @@ def build_vllm_args(model_source: str, adapter_path: str | None, temp_dir: Path)
                  "path": adapter_path, "base_model_name": base_name})]
     prompt = os.environ.get("TT_SYSTEM_PROMPT", "").strip()
     if prompt:
-        source = Path(model_source)
-        template_file = source / "chat_template.jinja"
-        if template_file.is_file():
-            template = template_file.read_text(encoding="utf-8")
-        else:
-            template = json.loads((source / "tokenizer_config.json").read_text())["chat_template"]
-        if not isinstance(template, str):
-            raise ValueError("Serving requires an unambiguous string chat template.")
+        template = native_chat_template(model_source)
+        from vllm.renderers.hf import _detect_content_format
+        content_format = _detect_content_format(template, default="string")
         # Literal data, never prompt text interpolated as Jinja source. Preserve
         # tool calls/results verbatim while merging leading system context.
         prefix = "{%- set tt = namespace(system=" + json.dumps(prompt) + ", history=[]) -%}"
@@ -101,7 +123,9 @@ def build_vllm_args(model_source: str, adapter_path: str | None, temp_dir: Path)
         prefix += "{%- set tt.system = tt.system + '\\n\\n' + part.text -%}"
         prefix += "{%- endif -%}{%- endfor -%}{%- endif -%}"
         prefix += "{%- else -%}{%- set tt.history = tt.history + [message] -%}{%- endif -%}{%- endfor -%}"
-        prefix += "{%- set messages = [{'role': 'system', 'content': tt.system}] + tt.history -%}"
+        # Match the ORIGINAL native template's expected content representation.
+        system_content = "[{'type': 'text', 'text': tt.system}]" if content_format == "openai" else "tt.system"
+        prefix += "{%- set messages = [{'role': 'system', 'content': " + system_content + "}] + tt.history -%}"
         path = temp_dir / "chat-template.jinja"
         path.write_text(prefix + template, encoding="utf-8")
         args += ["--chat-template", str(path)]
@@ -111,9 +135,14 @@ def build_vllm_args(model_source: str, adapter_path: str | None, temp_dir: Path)
 def prepare_model_source() -> str:
     base = os.environ["TT_BASE_MODEL"]
     revision = os.environ.get("TT_BASE_MODEL_REVISION")
-    assert_certified_base_model_revision(base, revision, "Serving base model revision")
-    configure_hugging_face_cache(os.environ.get("HF_HOME"))
     source = os.environ.get("TT_MODEL_SOURCE", base)
+    assert_certified_base_model_revision(
+        base,
+        revision,
+        "Serving base model revision",
+        local_source=source != base and Path(source).is_dir(),
+    )
+    configure_hugging_face_cache(os.environ.get("HF_HOME"))
     if not Path(source).is_dir():
         if source != base:
             raise ValueError("Explicit serving snapshot is not an existing directory.")
@@ -143,6 +172,7 @@ def run_server() -> None:
             model_source = prepare_model_source()
             adapter = resolve_adapter_path(os.environ.get("TT_MODEL_ARTIFACT"), Path(temp))
         args = build_vllm_args(model_source, adapter, Path(temp))
+        pin_native_content_format(args, model_source)
         from vllm.entrypoints.cli.main import main
         previous = sys.argv
         try:

package/training/adapter/src/train.py

@@ -27,6 +27,7 @@ from model_contract import (
     assert_certified_model_config,
     chat_template_kwargs,
     parse_lora_target_modules,
+    uncertified_load_kwargs,
 )
 from sft_data import IGNORE_INDEX, build_assistant_only_example
 
@@ -82,6 +83,7 @@ def assert_certified_request() -> None:
         str(requested_model),
         hp("base_model_revision"),
         "Training base model revision",
+        local_source=Path(resolve_model_source()).is_dir(),
     )
     loader = hp("model_loader", "causal_lm")
     if loader not in ("causal_lm", "image_text_to_text"):
@@ -199,6 +201,11 @@ def create_model_and_tokenizer(model_source: str) -> tuple[Any, Any, torch.dtype
         raise ValueError(f"{hp('base_model', CERTIFIED_BASE_MODEL)} tokenizer has no EOS token")
     if tokenizer.pad_token_id is None:
         tokenizer.pad_token = tokenizer.eos_token
+    if not tokenizer.chat_template:
+        raise ValueError(
+            f"{hp('base_model', CERTIFIED_BASE_MODEL)} has no chat template; "
+            "choose an instruct or chat variant of the model"
+        )
 
     dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
     # Multimodal checkpoints whose text tower is not registered in the
@@ -212,6 +219,7 @@ def create_model_and_tokenizer(model_source: str) -> tuple[Any, Any, torch.dtype
     model = model_class.from_pretrained(
         model_source,
         **source_kwargs,
+        **uncertified_load_kwargs(str(hp("base_model", CERTIFIED_BASE_MODEL))),
         dtype=dtype,
         device_map={"": torch.cuda.current_device()},
     )

package/training/decision/pyproject.toml

@@ -0,0 +1,25 @@
+[project]
+name = "tuned-tensor-decision"
+version = "0.1.0"
+requires-python = ">=3.11"
+dependencies = [
+  "huggingface-hub>=0.36.0",
+  "laya==0.3.22",
+  "numpy>=2.0.0",
+  "safetensors>=0.4.5",
+  "torch>=2.11.0",
+  "transformers>=4.48.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/decision/src/common.py

@@ -0,0 +1,127 @@
+"""Shared helpers for the typed decision-model engine (Laya-format checkpoints).
+
+Node owns the behavior spec, splits, scoring and reports. This package only
+loads a decision checkpoint, predicts labels for opaque IDs, and fine-tunes the
+checkpoint's encoder and decision head on labelled rows.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import tempfile
+from pathlib import Path
+from typing import Any
+
+import torch
+
+CHECKPOINT_FILES = ("rl_agent_config.json", "model.safetensors", "tokenizer/*", "encoder/*")
+QUESTION_ID = "decision"
+
+
+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 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 write_json(path: str | Path, value: Any) -> None:
+    write_text(path, json.dumps(value, indent=2) + "\n")
+
+
+def write_jsonl(path: str | Path, rows: list[Any]) -> None:
+    write_text(path, "".join(json.dumps(row, ensure_ascii=False) + "\n" for row in rows))
+
+
+def write_text(path: str | Path, content: str) -> None:
+    destination = Path(path)
+    ensure_private_directory(destination.parent)
+    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(content)
+            output.flush()
+            os.fsync(output.fileno())
+        os.replace(temporary, destination)
+    finally:
+        temporary.unlink(missing_ok=True)
+
+
+def read_jsonl(path: str | Path) -> list[dict[str, Any]]:
+    rows = []
+    for number, line in enumerate(Path(path).read_text(encoding="utf-8").splitlines(), start=1):
+        if not line.strip():
+            continue
+        row = json.loads(line)
+        if not isinstance(row, dict):
+            raise ValueError(f"{path}:{number} must be a JSON object")
+        rows.append(row)
+    return rows
+
+
+def pick_device(requested: str | None) -> torch.device:
+    if requested and requested != "auto":
+        device = torch.device(requested)
+        if device.type == "cuda" and not torch.cuda.is_available():
+            raise RuntimeError("Decision device cuda was requested, but PyTorch cannot see a CUDA GPU.")
+        if device.type == "mps" and not torch.backends.mps.is_available():
+            raise RuntimeError("Decision device mps was requested, but PyTorch cannot use Apple MPS.")
+        return device
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    if torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+def option_labels(question: dict[str, Any]) -> list[str]:
+    """Answer labels in the model's option order; these are the strings Node scores."""
+    kind = question["type"]
+    criteria = question.get("criteria")
+    if kind == "choice":
+        return [str(label) for label in (criteria.keys() if isinstance(criteria, dict) else criteria)]
+    if kind == "noul":
+        return ["false", "true"]
+    if kind == "score":
+        return [str(index) for index in range(len(criteria))]
+    raise ValueError(f"Unsupported decision type: {kind!r}")
+
+
+def label_index(question: dict[str, Any], output: str) -> int:
+    labels = option_labels(question)
+    normalized = output.strip().lower()
+    for index, label in enumerate(labels):
+        if label.strip().lower() == normalized:
+            return index
+    raise ValueError(f"Expected output {output!r} is not one of the decision labels {labels}")
+
+
+def resolve_checkpoint_dir(model: str, revision: str | None) -> Path:
+    """Return a local checkpoint directory, downloading only the files a load needs."""
+    local = Path(model)
+    if local.is_dir():
+        return local
+    if local.is_absolute() or model.startswith(("./", "../")):
+        raise FileNotFoundError(f"Decision checkpoint directory not found: {model}")
+    from huggingface_hub import snapshot_download
+
+    return Path(snapshot_download(model, revision=revision, allow_patterns=list(CHECKPOINT_FILES)))
+
+
+def load_agent(model: str, revision: str | None, device: torch.device):
+    from laya import Agent
+
+    checkpoint = resolve_checkpoint_dir(model, revision)
+    # Load the resolved directory so Hub and local checkpoints take the same path.
+    agent = Agent(str(checkpoint), device=str(device))
+    return agent, checkpoint

package/training/decision/src/evaluate.py

@@ -0,0 +1,70 @@
+"""Predict one typed decision per input row.
+
+The config names only opaque row IDs and inputs. Reference labels stay in Node,
+which joins these predictions to trusted expected outputs before scoring.
+"""
+from __future__ import annotations
+
+import time
+from typing import Any
+
+import numpy as np
+
+from common import QUESTION_ID, load_agent, load_config, option_labels, pick_device, read_jsonl, write_json, write_jsonl
+
+
+def answer_distribution(question: dict[str, Any], answer: dict[str, Any]) -> dict[str, float]:
+    labels = option_labels(question)
+    if question["type"] == "noul":
+        p_true = float(answer["noul"])
+        return {"false": 1.0 - p_true, "true": p_true}
+    # Laya keys choice probabilities by label and score probabilities by level index.
+    return {label: float(answer["probabilities"][label]) for label in labels}
+
+
+def predict_rows(agent, question: dict[str, Any], rows: list[dict[str, Any]], batch_size: int) -> list[dict[str, Any]]:
+    predictions: list[dict[str, Any]] = []
+    for start in range(0, len(rows), batch_size):
+        batch = rows[start:start + batch_size]
+        began = time.perf_counter()
+        results = agent.predict_batch([row["input"] for row in batch], {QUESTION_ID: question}, batch_size=batch_size)
+        latency_ms = int(round((time.perf_counter() - began) * 1000 / max(1, len(batch))))
+        for row, result in zip(batch, results):
+            distribution = answer_distribution(question, result["answers"][QUESTION_ID])
+            labels = list(distribution)
+            values = np.array([distribution[label] for label in labels])
+            predictions.append({
+                "id": row["id"],
+                "prediction": labels[int(values.argmax())],
+                "probabilities": {label: round(float(value), 6) for label, value in distribution.items()},
+                "confidence": round(float(values.max()), 6),
+                "latency_ms": latency_ms,
+            })
+    return predictions
+
+
+def main() -> None:
+    config = load_config()
+    device = pick_device(config.get("device"))
+    rows = read_jsonl(config["inputs_path"])
+    for row in rows:
+        if set(row) != {"id", "input"}:
+            raise ValueError("Decision evaluation rows must contain only id and input.")
+    began = time.perf_counter()
+    agent, checkpoint = load_agent(config["model"], config.get("revision"), device)
+    load_seconds = time.perf_counter() - began
+    predictions = predict_rows(agent, config["question"], rows, int(config.get("batch_size", 16)))
+    write_jsonl(config["output_path"], predictions)
+    write_json(config["metrics_path"], {
+        "ok": True,
+        "model": config["model"],
+        "revision": config.get("revision"),
+        "checkpoint_dir": str(checkpoint),
+        "device": str(agent.device),
+        "rows": len(predictions),
+        "load_seconds": round(load_seconds, 3),
+    })
+
+
+if __name__ == "__main__":
+    main()

package/training/decision/src/train.py

@@ -0,0 +1,181 @@
+"""Fine-tune a Laya-format decision checkpoint on labelled rows for one typed question.
+
+The encoder and decision head are trained with cross-entropy over the option
+markers, a strictly proper scoring rule, so the tuned probabilities remain a
+calibration target rather than only an argmax. Options are shuffled per example
+so the head learns option meaning instead of option position. The output
+directory is a complete checkpoint that `laya.load(<dir>)` accepts.
+"""
+from __future__ import annotations
+
+import json
+import random
+import shutil
+import time
+from contextlib import nullcontext
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from laya.common import QTYPES, build_sequence, collate_items
+from safetensors.torch import save_file
+
+from common import (
+    ensure_private_directory,
+    label_index,
+    load_agent,
+    load_config,
+    option_labels,
+    pick_device,
+    read_jsonl,
+    write_json,
+)
+
+
+def internal_question(question: dict[str, Any]) -> dict[str, Any]:
+    """Laya's internal question form; mirrors `Agent._to_internal` for the fields TT emits."""
+    criteria = question.get("criteria")
+    if question["type"] == "choice" and isinstance(criteria, list):
+        criteria = {label: None for label in criteria}
+    return {"t": question["type"], "ins": question["instructions"], "crit": criteria}
+
+
+def encode_example(tok, cfg: dict[str, Any], question: dict[str, Any], text: str, label: int, order: list[int]) -> dict[str, Any]:
+    ids, markers = build_sequence(
+        tok,
+        text,
+        question,
+        cfg.get("max_len", 512),
+        cfg.get("head_max_len", 192),
+        option_order=order,
+    )
+    if len(markers) != len(order):
+        raise ValueError("A decision option marker was truncated; shorten the instructions or option descriptions.")
+    # Slot s shows option order[s], so the target is the slot holding the labelled option.
+    return {"ids": ids, "markers": markers, "qtype": QTYPES[question["t"]], "label": order.index(label)}
+
+
+def autocast(device: torch.device):
+    if device.type == "cuda":
+        dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+        return torch.autocast("cuda", dtype=dtype)
+    return nullcontext()
+
+
+def save_checkpoint(model, source: Path, output: Path, cfg: dict[str, Any], provenance: dict[str, Any]) -> None:
+    ensure_private_directory(output)
+    for name in ("tokenizer", "encoder"):
+        if (source / name).is_dir():
+            shutil.copytree(source / name, output / name, dirs_exist_ok=True)
+    state = {key: value.detach().to("cpu").contiguous() for key, value in model.state_dict().items()}
+    # The base checkpoint's fitted temperatures describe its logits, not the tuned ones.
+    state["temperature"] = torch.ones_like(state["temperature"])
+    save_file(state, str(output / "model.safetensors"))
+    tuned = {key: value for key, value in cfg.items() if key != "temperature_by_options"}
+    tuned["temperature"] = [1.0, 1.0, 1.0]
+    tuned["tuned_tensor"] = provenance
+    (output / "rl_agent_config.json").write_text(json.dumps(tuned, indent=2) + "\n", encoding="utf-8")
+    # Copied tokenizer/encoder files keep the cache's modes; keep the whole checkpoint owner-only.
+    for path in output.rglob("*"):
+        path.chmod(0o700 if path.is_dir() else 0o600)
+
+
+def main() -> None:
+    config = load_config()
+    seed = int(config.get("seed", 0))
+    random.seed(seed)
+    torch.manual_seed(seed)
+    device = pick_device(config.get("device"))
+    question = config["question"]
+    labels = option_labels(question)
+    rows = read_jsonl(config["train_path"])
+    if not rows:
+        raise ValueError("Decision training needs at least one labelled row.")
+    targets = [label_index(question, str(row["output"])) for row in rows]
+
+    agent, checkpoint = load_agent(config["model"], config.get("revision"), device)
+    model = agent.model
+    if config.get("freeze_encoder", False):
+        for parameter in model.encoder.parameters():
+            parameter.requires_grad_(False)
+    trainable = [parameter for parameter in model.parameters() if parameter.requires_grad]
+    optimizer = torch.optim.AdamW(trainable, lr=float(config.get("learning_rate", 2e-5)), weight_decay=0.01)
+    epochs = int(config.get("epochs", 3))
+    batch_size = int(config.get("batch_size", 8))
+    shuffle_options = bool(config.get("shuffle_options", True))
+    steps_per_epoch = (len(rows) + batch_size - 1) // batch_size
+    total_steps = epochs * steps_per_epoch
+    warmup = max(1, total_steps // 10)
+    scheduler = torch.optim.lr_scheduler.LambdaLR(
+        optimizer,
+        lambda step: min(1.0, (step + 1) / warmup) * max(0.0, (total_steps - step) / max(1, total_steps - warmup + 1)),
+    )
+    q = internal_question(question)
+    pad_id = agent.tok.pad_token_id if agent.tok.pad_token_id is not None else 0
+
+    model.train()
+    history: list[dict[str, Any]] = []
+    began = time.perf_counter()
+    step = 0
+    for epoch in range(epochs):
+        order_of_rows = list(range(len(rows)))
+        random.shuffle(order_of_rows)
+        epoch_loss = 0.0
+        for start in range(0, len(rows), batch_size):
+            items = []
+            for index in order_of_rows[start:start + batch_size]:
+                order = list(range(len(labels)))
+                if shuffle_options:
+                    random.shuffle(order)
+                items.append([encode_example(agent.tok, agent.cfg, q, str(rows[index]["input"]), targets[index], order)])
+            batch = collate_items(items, pad_id)
+            with autocast(device):
+                logits, _ = model(
+                    batch["input_ids"].to(device),
+                    batch["attention_mask"].to(device),
+                    batch["marker_pos"].to(device),
+                    batch["marker_mask"].to(device),
+                    batch["qtype"].to(device),
+                )
+            loss = F.cross_entropy(logits.float(), batch["label"].to(device))
+            optimizer.zero_grad(set_to_none=True)
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(trainable, 1.0)
+            optimizer.step()
+            scheduler.step()
+            step += 1
+            epoch_loss += float(loss.item()) * len(items)
+            print(json.dumps({"event": "step", "epoch": epoch + 1, "step": step, "total_steps": total_steps, "loss": round(float(loss.item()), 6)}), flush=True)
+        history.append({"epoch": epoch + 1, "loss": round(epoch_loss / len(rows), 6)})
+
+    model.eval()
+    seconds = time.perf_counter() - began
+    provenance = {
+        "base_model": config["model"],
+        "base_revision": config.get("revision"),
+        "question_type": question["type"],
+        "labels": labels,
+        "train_rows": len(rows),
+        "epochs": epochs,
+        "learning_rate": float(config.get("learning_rate", 2e-5)),
+        "freeze_encoder": bool(config.get("freeze_encoder", False)),
+        "seed": seed,
+    }
+    output = Path(config["output_dir"])
+    save_checkpoint(model, checkpoint, output, agent.cfg, provenance)
+    write_json(config["metrics_path"], {
+        "ok": True,
+        "device": str(device),
+        "train_rows": len(rows),
+        "epochs": epochs,
+        "steps": step,
+        "train_seconds": round(seconds, 3),
+        "loss_history": history,
+        "final_loss": history[-1]["loss"],
+        "trainable_parameters": sum(parameter.numel() for parameter in trainable),
+    })
+
+
+if __name__ == "__main__":
+    main()