@tuned-tensor/cli 0.9.0 → 0.9.1

package/training/local-runner/src/serve.py

@@ -0,0 +1,272 @@
+from __future__ import annotations
+
+import hmac
+import json
+import math
+import os
+import sys
+import threading
+import time
+import uuid
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Any
+
+from evaluate import (
+    configure_hugging_face_cache,
+    import_runtime_dependencies,
+    load_text_model,
+    resolve_adapter_path,
+    sampling_kwargs,
+)
+
+
+MODEL_ARTIFACT = os.environ.get("TT_MODEL_ARTIFACT")
+BASE_MODEL = os.environ["TT_BASE_MODEL"]
+BASE_MODEL_REVISION = os.environ.get("TT_BASE_MODEL_REVISION")
+MODEL_NAME = os.environ.get("TT_MODEL_NAME", "tuned-tensor-local")
+MODEL_LOADER = os.environ.get("TT_MODEL_LOADER", "causal_lm")
+SYSTEM_PROMPT = os.environ.get("TT_SYSTEM_PROMPT", "").strip()
+HOST = os.environ.get("TT_HOST", "127.0.0.1")
+PORT = int(os.environ.get("TT_PORT", "8000"))
+DEVICE_REQUEST = os.environ.get("TT_DEVICE", "cuda")
+DEFAULT_MAX_TOKENS = int(os.environ.get("TT_MAX_TOKENS", "512"))
+DEFAULT_TEMPERATURE = float(os.environ.get("TT_TEMPERATURE", "0"))
+DEFAULT_TOP_P = float(os.environ.get("TT_TOP_P", "1"))
+MAX_REQUEST_BYTES = 2 * 1024 * 1024
+MAX_PROMPT_CHARS = 100_000
+MAX_PROMPT_TOKENS = 16_384
+MAX_MESSAGES = 128
+API_KEY = os.environ.get("TT_API_KEY", "")
+MAX_CONCURRENT_REQUESTS = int(os.environ.get("TT_MAX_CONCURRENT_REQUESTS", "1"))
+
+
+if MODEL_LOADER != "causal_lm":
+    raise ValueError("The bundled model server is text-only and requires TT_MODEL_LOADER=causal_lm")
+configure_hugging_face_cache(os.environ.get("HF_HOME"))
+import_runtime_dependencies()
+TEMP_DIR = TemporaryDirectory(prefix="tt-local-serve-")
+ADAPTER_PATH = (
+    resolve_adapter_path(MODEL_ARTIFACT, Path(TEMP_DIR.name))
+    if MODEL_ARTIFACT
+    else None
+)
+MODEL_PAYLOAD = {
+    "base_model": BASE_MODEL,
+    "base_model_revision": BASE_MODEL_REVISION,
+    "device": DEVICE_REQUEST,
+    "model_loader": MODEL_LOADER,
+}
+MODEL, TOKENIZER, DEVICE = load_text_model(MODEL_PAYLOAD, ADAPTER_PATH)
+GENERATION_LOCK = threading.Lock()
+REQUEST_SLOTS = threading.BoundedSemaphore(MAX_CONCURRENT_REQUESTS)
+
+
+def text_content(content: Any) -> str:
+    if isinstance(content, str):
+        return content
+    if content is None:
+        return ""
+    if not isinstance(content, list):
+        raise ValueError("Message content must be text or an array of text parts.")
+    parts: list[str] = []
+    for part in content:
+        if not isinstance(part, dict) or part.get("type") != "text":
+            raise ValueError("The bundled Qwen model server accepts text content only.")
+        text = part.get("text")
+        if not isinstance(text, str):
+            raise ValueError("Every text content part must contain a string text field.")
+        parts.append(text)
+    return "\n".join(parts)
+
+
+def normalize_messages(raw_messages: Any) -> list[dict[str, str]]:
+    if not isinstance(raw_messages, list) or not raw_messages:
+        raise ValueError("Request must include a non-empty messages array.")
+    if len(raw_messages) > MAX_MESSAGES:
+        raise ValueError(f"Request exceeds the {MAX_MESSAGES}-message limit.")
+
+    system_parts = [SYSTEM_PROMPT] if SYSTEM_PROMPT else []
+    conversation: list[dict[str, str]] = []
+    for raw in raw_messages:
+        if not isinstance(raw, dict) or raw.get("role") not in {"system", "user", "assistant", "tool"}:
+            raise ValueError("Each message must have a supported role and text content.")
+        role = str(raw["role"])
+        content = text_content(raw.get("content", ""))
+        if role == "system":
+            if content.strip():
+                system_parts.append(content.strip())
+        else:
+            conversation.append({"role": role, "content": content})
+
+    if not conversation:
+        raise ValueError("Request must contain at least one non-system message.")
+    messages: list[dict[str, str]] = []
+    if system_parts:
+        # Qwen permits one leading system message. Merge the invariant owner
+        # prompt and any client context instead of triggering an implicit
+        # template fallback with duplicate system turns.
+        messages.append({"role": "system", "content": "\n\n".join(system_parts)})
+    messages.extend(conversation)
+    if sum(len(message["content"]) for message in messages) > MAX_PROMPT_CHARS:
+        raise ValueError(f"Prompt exceeds the {MAX_PROMPT_CHARS}-character limit.")
+    return messages
+
+
+def generate_text(messages: list[dict[str, str]], generation: dict[str, Any]) -> tuple[str, int, int]:
+    try:
+        prompt = TOKENIZER.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+    except Exception as exc:
+        raise ValueError(f"Qwen chat-template rendering failed: {exc}") from exc
+    inputs = TOKENIZER(prompt, return_tensors="pt")
+    if int(inputs["input_ids"].shape[-1]) > MAX_PROMPT_TOKENS:
+        raise ValueError(f"Prompt exceeds the {MAX_PROMPT_TOKENS}-token limit.")
+    target_device = next(MODEL.parameters()).device
+    inputs = {key: value.to(target_device) for key, value in inputs.items()}
+    with GENERATION_LOCK:
+        import torch
+
+        with torch.inference_mode():
+            generated = MODEL.generate(
+                **inputs,
+                max_new_tokens=int(generation["max_new_tokens"]),
+                pad_token_id=TOKENIZER.eos_token_id,
+                **sampling_kwargs(generation),
+            )
+    prompt_tokens = int(inputs["input_ids"].shape[-1])
+    completion_tokens = generated[0][prompt_tokens:]
+    content = TOKENIZER.decode(completion_tokens, skip_special_tokens=True).strip()
+    return content, prompt_tokens, int(completion_tokens.shape[-1])
+
+
+def bounded_number(value: Any, default: float, minimum: float, maximum: float) -> float:
+    number = float(default if value is None else value)
+    if not math.isfinite(number) or number < minimum or number > maximum:
+        raise ValueError(f"Generation value must be between {minimum} and {maximum}.")
+    return number
+
+
+def bounded_integer(value: Any, default: int, minimum: int, maximum: int) -> int:
+    number = bounded_number(value, default, minimum, maximum)
+    if not number.is_integer():
+        raise ValueError(f"Generation value must be an integer between {minimum} and {maximum}.")
+    return int(number)
+
+
+class Handler(BaseHTTPRequestHandler):
+    server_version = "TunedTensorLocalServer/0.1"
+
+    def setup(self) -> None:
+        super().setup()
+        self.connection.settimeout(30)
+
+    def authorized(self) -> bool:
+        if not API_KEY:
+            return True
+        supplied = self.headers.get("authorization", "")
+        return hmac.compare_digest(supplied, "Bearer " + API_KEY)
+
+    def send_json(self, status: int, payload: Any) -> None:
+        body = json.dumps(payload).encode("utf-8")
+        self.send_response(status)
+        self.send_header("Content-Type", "application/json")
+        self.send_header("Content-Length", str(len(body)))
+        self.end_headers()
+        self.wfile.write(body)
+
+    def do_GET(self) -> None:  # noqa: N802
+        if not self.authorized():
+            self.send_json(401, {"error": {"message": "Unauthorized"}})
+            return
+        if self.path == "/health":
+            self.send_json(200, {"status": "ok", "model": MODEL_NAME, "device": DEVICE})
+            return
+        if self.path == "/v1/models":
+            self.send_json(200, {
+                "object": "list",
+                "data": [{"id": MODEL_NAME, "object": "model", "owned_by": "tuned-tensor-local"}],
+            })
+            return
+        self.send_json(404, {"error": {"message": "Not found"}})
+
+    def do_POST(self) -> None:  # noqa: N802
+        if not self.authorized():
+            self.send_json(401, {"error": {"message": "Unauthorized"}})
+            return
+        if self.path not in {"/v1/chat/completions", "/chat/completions"}:
+            self.send_json(404, {"error": {"message": "Not found"}})
+            return
+        try:
+            length = int(self.headers.get("content-length", "0"))
+            if length <= 0 or length > MAX_REQUEST_BYTES:
+                raise ValueError("Request body is empty or too large.")
+            body = json.loads(self.rfile.read(length))
+            if not isinstance(body, dict):
+                raise ValueError("Request body must be a JSON object.")
+            if body.get("stream"):
+                self.send_json(400, {"error": {"message": "Streaming is not supported yet."}})
+                return
+            messages = normalize_messages(body.get("messages"))
+            generation = {
+                "max_new_tokens": bounded_integer(body.get("max_tokens"), DEFAULT_MAX_TOKENS, 1, 8192),
+                "temperature": bounded_number(body.get("temperature"), DEFAULT_TEMPERATURE, 0, 5),
+                "top_p": bounded_number(body.get("top_p"), DEFAULT_TOP_P, 0, 1),
+            }
+            if not REQUEST_SLOTS.acquire(blocking=False):
+                self.send_json(429, {"error": {"message": "The local model is busy; retry shortly."}})
+                return
+            started = time.perf_counter()
+            try:
+                content, prompt_tokens, completion_tokens = generate_text(messages, generation)
+            finally:
+                REQUEST_SLOTS.release()
+            latency_ms = round((time.perf_counter() - started) * 1000)
+            self.send_json(200, {
+                "id": "chatcmpl-" + uuid.uuid4().hex,
+                "object": "chat.completion",
+                "created": int(time.time()),
+                "model": MODEL_NAME,
+                "choices": [{
+                    "index": 0,
+                    "message": {"role": "assistant", "content": content},
+                    "finish_reason": "stop",
+                }],
+                "usage": {
+                    "prompt_tokens": prompt_tokens,
+                    "completion_tokens": completion_tokens,
+                    "total_tokens": prompt_tokens + completion_tokens,
+                },
+                "tt_local": {"latency_ms": latency_ms, "device": DEVICE},
+            })
+        except (ValueError, TypeError, json.JSONDecodeError) as exc:
+            self.send_json(400, {"error": {"message": str(exc)}})
+        except Exception as exc:
+            print(f"Model server error: {type(exc).__name__}: {exc}", file=sys.stderr, flush=True)
+            self.send_json(500, {"error": {"message": "Internal model server error."}})
+
+    def log_message(self, fmt: str, *args: Any) -> None:
+        print("%s - %s" % (self.address_string(), fmt % args), flush=True)
+
+
+def main() -> None:
+    print(f"Serving {MODEL_NAME} on http://{HOST}:{PORT}", flush=True)
+    print(f"OpenAI-compatible endpoint: http://{HOST}:{PORT}/v1/chat/completions", flush=True)
+    print(f"Device: {DEVICE}", flush=True)
+    server = ThreadingHTTPServer((HOST, PORT), Handler)
+    server.daemon_threads = True
+    server.request_queue_size = max(2, MAX_CONCURRENT_REQUESTS * 2)
+    try:
+        server.serve_forever()
+    except KeyboardInterrupt:
+        print("Model server stopped.", flush=True)
+    finally:
+        server.server_close()
+
+
+if __name__ == "__main__":
+    main()

package/training/local-runner/src/sft_data.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+
+IGNORE_INDEX = -100
+
+
+def _token_ids(value: Any, description: str) -> list[int]:
+    if isinstance(value, Mapping):
+        value = value.get("input_ids")
+    if not isinstance(value, Sequence) or isinstance(value, (str, bytes)):
+        raise TypeError(f"{description} did not return a token-id sequence")
+    token_ids = list(value)
+    if not token_ids or any(not isinstance(token_id, int) for token_id in token_ids):
+        raise ValueError(f"{description} returned empty or invalid token ids")
+    return token_ids
+
+
+def build_assistant_only_example(
+    tokenizer: Any,
+    messages: Any,
+    *,
+    max_length: int,
+) -> dict[str, list[int]]:
+    """Tokenize one chat row while computing loss only on the final assistant turn.
+
+    The entire assistant completion is reserved before the prompt is truncated.
+    Prompt truncation therefore removes the oldest prompt tokens and can never
+    silently discard the expected answer.
+    """
+    if max_length < 2:
+        raise ValueError("max_length must be at least 2")
+    if not isinstance(messages, list) or len(messages) < 2:
+        raise ValueError("messages must contain prompt context and a final assistant turn")
+    if any(not isinstance(message, dict) for message in messages):
+        raise ValueError("every message must be an object")
+    allowed_roles = {"system", "user", "assistant", "tool"}
+    for message in messages:
+        role = message.get("role")
+        if role not in allowed_roles:
+            raise ValueError(f"unsupported message role: {role!r}")
+        if not isinstance(message.get("content"), str):
+            raise ValueError("the bundled trainer accepts text-only message content")
+
+    assistant = messages[-1]
+    if assistant.get("role") != "assistant":
+        raise ValueError("the final message must have role=assistant")
+    content = assistant.get("content")
+    if not isinstance(content, str) or not content.strip():
+        raise ValueError("the final assistant message must contain non-empty text")
+
+    try:
+        prompt_ids = _token_ids(
+            tokenizer.apply_chat_template(
+                messages[:-1],
+                tokenize=True,
+                add_generation_prompt=True,
+                return_dict=False,
+            ),
+            "prompt chat template",
+        )
+        full_ids = _token_ids(
+            tokenizer.apply_chat_template(
+                messages,
+                tokenize=True,
+                add_generation_prompt=False,
+                return_dict=False,
+            ),
+            "full chat template",
+        )
+    except Exception as exc:
+        raise ValueError(f"Qwen chat-template rendering failed: {exc}") from exc
+
+    if full_ids[: len(prompt_ids)] != prompt_ids:
+        raise ValueError(
+            "Qwen chat-template output is not prefix-aligned; refusing to guess the assistant loss boundary"
+        )
+    completion_ids = full_ids[len(prompt_ids) :]
+    if not completion_ids:
+        raise ValueError("Qwen chat template produced no assistant completion tokens")
+    if len(completion_ids) >= max_length:
+        raise ValueError(
+            f"assistant completion requires {len(completion_ids)} tokens, but max_length={max_length}; "
+            "increase max_seq_length so the answer is not truncated"
+        )
+
+    prompt_budget = max_length - len(completion_ids)
+    retained_prompt_ids = prompt_ids[-prompt_budget:]
+    input_ids = [*retained_prompt_ids, *completion_ids]
+    labels = [IGNORE_INDEX] * len(retained_prompt_ids) + completion_ids.copy()
+    return {
+        "input_ids": input_ids,
+        "attention_mask": [1] * len(input_ids),
+        "labels": labels,
+    }

package/training/local-runner/src/train.py

@@ -0,0 +1,289 @@
+from __future__ import annotations
+
+import json
+import os
+import tarfile
+import time
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Any
+
+import torch
+from peft import LoraConfig, get_peft_model
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.data import Dataset
+from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
+
+from model_contract import (
+    CERTIFIED_BASE_MODEL,
+    assert_certified_model_config,
+)
+from sft_data import IGNORE_INDEX, build_assistant_only_example
+
+
+TRAINING_DIR = Path(os.environ.get("SM_CHANNEL_TRAINING", "/opt/ml/input/data/training"))
+BASE_MODEL_DIR = Path(os.environ.get("SM_CHANNEL_BASE_MODEL", "/opt/ml/input/data/base_model"))
+HYPERPARAMETERS_PATH = Path(
+    os.environ.get("TT_HYPERPARAMETERS_PATH", "/opt/ml/input/config/hyperparameters.json")
+)
+MODEL_DIR = Path(os.environ.get("SM_MODEL_DIR", "/opt/ml/model"))
+OUTPUT_DIR = Path(os.environ.get("SM_OUTPUT_DIR", "/opt/ml/output"))
+
+
+def load_hyperparameters() -> dict[str, str]:
+    if not HYPERPARAMETERS_PATH.is_file():
+        return {}
+    raw = json.loads(HYPERPARAMETERS_PATH.read_text(encoding="utf-8"))
+    if not isinstance(raw, dict):
+        raise ValueError("Hyperparameters must be a JSON object")
+    return {str(key): str(value) for key, value in raw.items()}
+
+
+HP = load_hyperparameters()
+
+
+def hp(name: str, default: str | None = None) -> str | None:
+    value = os.getenv(f"SM_HP_{name.upper()}", HP.get(name, default))
+    if value is None:
+        return None
+    value = str(value).strip()
+    if len(value) >= 2 and value[0] == value[-1] == '"':
+        return value[1:-1]
+    return value
+
+
+def hp_int(name: str, default: int) -> int:
+    return int(hp(name, str(default)) or default)
+
+
+def hp_float(name: str, default: float) -> float:
+    return float(hp(name, str(default)) or default)
+
+
+def model_revision_kwargs(model_source: str) -> dict[str, str]:
+    revision = hp("base_model_revision")
+    return {"revision": revision} if revision and not Path(model_source).exists() else {}
+
+
+def assert_certified_request() -> None:
+    requested_model = hp("base_model", CERTIFIED_BASE_MODEL)
+    if requested_model != CERTIFIED_BASE_MODEL:
+        raise ValueError(
+            f"The bundled trainer currently certifies only {CERTIFIED_BASE_MODEL}; got {requested_model!r}"
+        )
+    loader = hp("model_loader", "causal_lm")
+    if loader != "causal_lm":
+        raise ValueError(f"The bundled trainer is text-only and requires model_loader=causal_lm; got {loader!r}")
+def require_cuda() -> dict[str, Any]:
+    if not torch.cuda.is_available():
+        raise RuntimeError(
+            "TT Local bundled training requires an NVIDIA CUDA GPU. "
+            "Run this workflow on DGX Spark and verify it first with tt-local doctor."
+        )
+    device_index = torch.cuda.current_device()
+    properties = torch.cuda.get_device_properties(device_index)
+    return {
+        "device_index": device_index,
+        "device_name": properties.name,
+        "compute_capability": f"{properties.major}.{properties.minor}",
+        "bf16_supported": torch.cuda.is_bf16_supported(),
+    }
+
+
+def load_rows() -> list[dict[str, Any]]:
+    rows: list[dict[str, Any]] = []
+    for path in sorted(TRAINING_DIR.rglob("*.jsonl")):
+        for line_number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1):
+            if not line.strip():
+                continue
+            try:
+                row = json.loads(line)
+            except json.JSONDecodeError as exc:
+                raise ValueError(f"{path}:{line_number} contains invalid JSON") from exc
+            if not isinstance(row, dict) or "messages" not in row:
+                raise ValueError(f"{path}:{line_number} must contain a messages array")
+            rows.append(row)
+    if not rows:
+        raise ValueError(f"No chat JSONL rows found under {TRAINING_DIR}")
+    return rows
+
+
+def resolve_model_source() -> str:
+    if BASE_MODEL_DIR.is_dir() and any(BASE_MODEL_DIR.iterdir()):
+        return str(BASE_MODEL_DIR)
+    return CERTIFIED_BASE_MODEL
+
+
+class AssistantOnlyDataset(Dataset[dict[str, torch.Tensor]]):
+    def __init__(
+        self,
+        rows: list[dict[str, Any]],
+        tokenizer: Any,
+        *,
+        max_length: int,
+    ):
+        self.examples: list[dict[str, torch.Tensor]] = []
+        for row_index, row in enumerate(rows, start=1):
+            try:
+                tokenized = build_assistant_only_example(
+                    tokenizer,
+                    row["messages"],
+                    max_length=max_length,
+                )
+            except Exception as exc:
+                raise ValueError(f"Training row {row_index} is not valid Qwen SFT data: {exc}") from exc
+            self.examples.append({
+                key: torch.tensor(value, dtype=torch.long)
+                for key, value in tokenized.items()
+            })
+
+    def __len__(self) -> int:
+        return len(self.examples)
+
+    def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
+        return self.examples[index]
+
+
+class AssistantOnlyCollator:
+    def __init__(self, pad_token_id: int):
+        self.pad_token_id = pad_token_id
+
+    def __call__(self, features: list[dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
+        return {
+            "input_ids": pad_sequence(
+                [feature["input_ids"] for feature in features],
+                batch_first=True,
+                padding_value=self.pad_token_id,
+            ),
+            "attention_mask": pad_sequence(
+                [feature["attention_mask"] for feature in features],
+                batch_first=True,
+                padding_value=0,
+            ),
+            "labels": pad_sequence(
+                [feature["labels"] for feature in features],
+                batch_first=True,
+                padding_value=IGNORE_INDEX,
+            ),
+        }
+
+
+def create_model_and_tokenizer(model_source: str) -> tuple[Any, Any, torch.dtype]:
+    token = os.getenv("HF_TOKEN")
+    source_kwargs: dict[str, Any] = {
+        **model_revision_kwargs(model_source),
+        "trust_remote_code": False,
+        "token": token,
+    }
+    if not Path(model_source).exists():
+        source_kwargs["local_files_only"] = True
+    tokenizer = AutoTokenizer.from_pretrained(model_source, **source_kwargs)
+    if tokenizer.eos_token_id is None:
+        raise ValueError(f"{CERTIFIED_BASE_MODEL} tokenizer has no EOS token")
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+    model = AutoModelForCausalLM.from_pretrained(
+        model_source,
+        **source_kwargs,
+        dtype=dtype,
+        device_map={"": torch.cuda.current_device()},
+    )
+    assert_certified_model_config(model.config)
+    model.config.use_cache = False
+    return model, tokenizer, dtype
+
+
+def apply_lora(model: Any) -> Any:
+    config = LoraConfig(
+        r=hp_int("lora_rank", 16),
+        lora_alpha=hp_int("lora_alpha", 32),
+        lora_dropout=hp_float("lora_dropout", 0.05),
+        bias="none",
+        task_type="CAUSAL_LM",
+        target_modules="all-linear",
+    )
+    return get_peft_model(model, config)
+
+
+def create_model_archive() -> Path:
+    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+    archive_path = OUTPUT_DIR / "model.tar.gz"
+    with tarfile.open(archive_path, "w:gz") as archive:
+        archive.add(MODEL_DIR, arcname="model")
+    return archive_path
+
+
+def run_training(rows: list[dict[str, Any]], model_source: str) -> tuple[dict[str, Any], torch.dtype]:
+    model, tokenizer, dtype = create_model_and_tokenizer(model_source)
+    model = apply_lora(model)
+    dataset = AssistantOnlyDataset(
+        rows,
+        tokenizer,
+        max_length=hp_int("max_seq_length", 2048),
+    )
+    collator = AssistantOnlyCollator(tokenizer.pad_token_id)
+
+    with TemporaryDirectory(prefix="tt-local-sft-") as output_dir:
+        args = TrainingArguments(
+            output_dir=output_dir,
+            num_train_epochs=hp_int("n_epochs", 3),
+            learning_rate=hp_float("learning_rate", 0.00001),
+            per_device_train_batch_size=hp_int("per_device_train_batch_size", 1),
+            gradient_accumulation_steps=hp_int("gradient_accumulation_steps", 8),
+            logging_steps=1,
+            save_strategy="no",
+            report_to=[],
+            remove_unused_columns=False,
+            bf16=dtype == torch.bfloat16,
+            fp16=dtype == torch.float16,
+            optim="adamw_torch",
+        )
+        trainer = Trainer(
+            model=model,
+            args=args,
+            train_dataset=dataset,
+            data_collator=collator,
+        )
+        result = trainer.train()
+
+    model.save_pretrained(MODEL_DIR)
+    tokenizer.save_pretrained(MODEL_DIR)
+    return result.metrics, dtype
+
+
+def main() -> None:
+    started = time.time()
+    assert_certified_request()
+    cuda = require_cuda()
+    MODEL_DIR.mkdir(parents=True, exist_ok=True)
+    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+
+    rows = load_rows()
+    model_source = resolve_model_source()
+    metrics, dtype = run_training(rows, model_source)
+    archive_path = create_model_archive()
+    output_metrics = {
+        "training_rows": len(rows),
+        "base_model": CERTIFIED_BASE_MODEL,
+        "base_model_revision": hp("base_model_revision"),
+        "model_source": model_source,
+        "model_loader": "causal_lm",
+        "loss_mask": "final_assistant_only",
+        "device": "cuda",
+        "cuda_device": cuda["device_name"],
+        "cuda_compute_capability": cuda["compute_capability"],
+        "dtype": str(dtype).removeprefix("torch."),
+        "train_runtime": round(time.time() - started, 3),
+        **{key: float(value) for key, value in metrics.items() if isinstance(value, (int, float))},
+        "model_archive": str(archive_path),
+    }
+    metrics_json = json.dumps(output_metrics, indent=2)
+    (MODEL_DIR / "training-metrics.json").write_text(metrics_json, encoding="utf-8")
+    (OUTPUT_DIR / "training-metrics.json").write_text(metrics_json, encoding="utf-8")
+    print(metrics_json)
+
+
+if __name__ == "__main__":
+    main()