@tuned-tensor/cli 0.9.0 → 0.9.2

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

@@ -0,0 +1,274 @@
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import tarfile
+import time
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Any
+
+from model_contract import (
+    CERTIFIED_BASE_MODEL,
+    assert_certified_model_config,
+)
+
+MAX_ARCHIVE_MEMBERS = 20_000
+MAX_ARCHIVE_EXPANDED_BYTES = 20 * 1024 * 1024 * 1024
+
+
+def load_json(path: Path) -> dict[str, Any]:
+    value = json.loads(path.read_text(encoding="utf-8"))
+    if not isinstance(value, dict):
+        raise ValueError(f"Expected a JSON object in {path}")
+    return value
+
+
+def configure_hugging_face_cache(cache_home: str | None) -> None:
+    """Treat model_cache as HF_HOME and keep legacy overrides consistent."""
+    if not cache_home:
+        return
+    home = Path(cache_home).expanduser().absolute()
+    hub = home / "hub"
+    os.environ.update({
+        "HF_HOME": str(home),
+        "HF_HUB_CACHE": str(hub),
+        "HUGGINGFACE_HUB_CACHE": str(hub),
+    })
+    for deprecated in (
+        "TRANSFORMERS_CACHE",
+        "PYTORCH_TRANSFORMERS_CACHE",
+        "PYTORCH_PRETRAINED_BERT_CACHE",
+    ):
+        os.environ.pop(deprecated, None)
+
+
+def import_runtime_dependencies() -> None:
+    """Import libraries only after Hugging Face cache settings are final."""
+    global torch, PeftModel, AutoModelForCausalLM, AutoTokenizer
+
+    import torch as torch_module
+    from peft import PeftModel as peft_model
+    from transformers import (
+        AutoModelForCausalLM as auto_model_for_causal_lm,
+        AutoTokenizer as auto_tokenizer,
+    )
+
+    torch = torch_module
+    PeftModel = peft_model
+    AutoModelForCausalLM = auto_model_for_causal_lm
+    AutoTokenizer = auto_tokenizer
+
+
+def strip_file_uri(value: str | None) -> str | None:
+    if not value:
+        return value
+    if value.startswith("file://"):
+        return value[7:]
+    return value
+
+
+def _extract_adapter_archive(path: Path, destination: Path) -> None:
+    destination.mkdir(parents=True, exist_ok=True)
+    destination_root = destination.resolve()
+    with tarfile.open(path, "r:gz") as archive:
+        members = archive.getmembers()
+        if len(members) > MAX_ARCHIVE_MEMBERS:
+            raise ValueError(f"Model archive exceeds {MAX_ARCHIVE_MEMBERS} members")
+        expanded_bytes = sum(max(0, member.size) for member in members if member.isfile())
+        if expanded_bytes > MAX_ARCHIVE_EXPANDED_BYTES:
+            raise ValueError("Model archive exceeds the 20 GiB expanded-size limit")
+        for member in members:
+            member_path = (destination / member.name).resolve()
+            try:
+                member_path.relative_to(destination_root)
+            except ValueError as exc:
+                raise ValueError(f"Unsafe archive member: {member.name}") from exc
+            if member.issym() or member.islnk() or member.isdev():
+                raise ValueError(f"Unsafe archive member type: {member.name}")
+        archive.extractall(destination)
+
+
+def resolve_adapter_path(value: str | None, tmp: Path) -> str | None:
+    path_value = strip_file_uri(value)
+    if not path_value:
+        return None
+    path = Path(path_value)
+    if path.is_file() and path.name.endswith(".tar.gz"):
+        extracted = tmp / "adapter"
+        _extract_adapter_archive(path, extracted)
+        candidates = sorted(extracted.rglob("adapter_config.json"))
+        if candidates:
+            return str(candidates[0].parent)
+        raise ValueError(f"Adapter archive contains no adapter_config.json: {path}")
+    if not path.is_dir():
+        raise ValueError(f"Adapter path is not a directory or .tar.gz archive: {path}")
+    return str(path)
+
+
+def format_prompt(
+    tokenizer: Any,
+    system: str,
+    prompt: str,
+) -> str:
+    messages = [
+        {"role": "system", "content": system},
+        {"role": "user", "content": prompt},
+    ]
+    try:
+        return 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
+
+
+def resolve_device(device: str) -> str:
+    if device not in {"cuda", "cpu"}:
+        raise ValueError(f"device must be cuda or cpu; got {device!r}")
+    if device == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError("CUDA evaluation was requested, but torch.cuda.is_available() is false")
+    return device
+
+
+def _assert_certified_model_source(base_model: str) -> None:
+    if not Path(base_model).exists() and base_model != CERTIFIED_BASE_MODEL:
+        raise ValueError(
+            f"The bundled evaluator currently certifies only {CERTIFIED_BASE_MODEL}; got {base_model!r}"
+        )
+
+
+def load_text_model(payload: dict[str, Any], adapter_path: str | None):
+    base_model = str(payload["base_model"])
+    _assert_certified_model_source(base_model)
+    device = resolve_device(str(payload.get("device", "cuda")))
+    token = os.getenv("HF_TOKEN")
+    revision = payload.get("base_model_revision")
+    source_kwargs: dict[str, Any] = {
+        "trust_remote_code": False,
+        "token": token,
+    }
+    if not Path(base_model).exists():
+        source_kwargs["local_files_only"] = True
+    if revision and not Path(base_model).exists():
+        source_kwargs["revision"] = revision
+
+    tokenizer = AutoTokenizer.from_pretrained(base_model, **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 = None
+    if device == "cuda":
+        dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+
+    model_kwargs: dict[str, Any] = dict(source_kwargs)
+    if dtype is not None:
+        model_kwargs["dtype"] = dtype
+    if device == "cuda":
+        model_kwargs["device_map"] = {"": torch.cuda.current_device()}
+
+    model = AutoModelForCausalLM.from_pretrained(base_model, **model_kwargs)
+    assert_certified_model_config(model.config)
+    if adapter_path:
+        model = PeftModel.from_pretrained(model, adapter_path)
+    if device != "cuda":
+        model = model.to(device)
+    model.eval()
+    return model, tokenizer, device
+
+
+def sampling_kwargs(generation: dict[str, Any]) -> dict[str, Any]:
+    temperature = float(generation.get("temperature", 0))
+    if temperature <= 0:
+        return {"do_sample": False}
+    return {
+        "do_sample": True,
+        "temperature": temperature,
+        "top_p": float(generation.get("top_p", 1)),
+    }
+
+
+def generate_text_one(
+    model: Any,
+    tokenizer: Any,
+    system: str,
+    example: dict[str, Any],
+    generation: dict[str, Any],
+) -> dict[str, Any]:
+    if example.get("input_assets"):
+        raise ValueError("The bundled Qwen SFT evaluator is text-only and does not accept input_assets")
+    prompt = str(example["input"])
+    formatted = format_prompt(tokenizer, system, prompt)
+    inputs = tokenizer(formatted, return_tensors="pt")
+    target_device = next(model.parameters()).device
+    inputs = {key: value.to(target_device) for key, value in inputs.items()}
+    started = time.perf_counter()
+    with torch.inference_mode():
+        generated = model.generate(
+            **inputs,
+            max_new_tokens=int(generation.get("max_new_tokens", 256)),
+            pad_token_id=tokenizer.eos_token_id,
+            **sampling_kwargs(generation),
+        )
+    latency_ms = max(0, round((time.perf_counter() - started) * 1000))
+    input_length = int(inputs["input_ids"].shape[-1])
+    output_tokens = generated[0][input_length:]
+    actual = tokenizer.decode(output_tokens, skip_special_tokens=True).strip()
+    return {
+        "id": str(example["id"]),
+        "actual": actual,
+        "latency_ms": latency_ms,
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input", required=True)
+    parser.add_argument("--output", required=True)
+    args = parser.parse_args()
+
+    payload = load_json(Path(args.input))
+    if payload.get("protocol_version") != 2:
+        raise ValueError("Unsupported inference protocol; expected protocol_version 2")
+    if payload.get("model_loader") not in (None, "causal_lm"):
+        raise ValueError("The bundled evaluator is text-only and requires model_loader=causal_lm")
+    configure_hugging_face_cache(payload.get("model_cache"))
+    import_runtime_dependencies()
+
+    with TemporaryDirectory(prefix="tt-local-evaluate-") as tmp_dir:
+        adapter_path = resolve_adapter_path(payload.get("adapter_path"), Path(tmp_dir))
+        generation = payload.get("generation", {})
+        if not isinstance(generation, dict):
+            raise ValueError("generation must be a JSON object")
+        model, tokenizer, _device = load_text_model(payload, adapter_path)
+        results = [
+            generate_text_one(
+                model,
+                tokenizer,
+                str(payload.get("system", "")),
+                example,
+                generation,
+            )
+            for example in payload.get("examples", [])
+        ]
+
+    output = {
+        "provider": "transformers",
+        "model_id": payload.get("model_id"),
+        "base_model": payload.get("base_model"),
+        "adapter_path": payload.get("adapter_path"),
+        "generation_config": generation,
+        "results": results,
+    }
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    output_path.write_text(json.dumps(output, indent=2), encoding="utf-8")
+
+
+if __name__ == "__main__":
+    main()

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

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from typing import Any
+
+
+CERTIFIED_BASE_MODEL = "Qwen/Qwen3.5-2B"
+CERTIFIED_TEXT_CONFIG = {
+    "model_type": "qwen3_5_text",
+    "hidden_size": 2048,
+    "num_hidden_layers": 24,
+    "num_attention_heads": 8,
+    "num_key_value_heads": 2,
+    "intermediate_size": 6144,
+    "vocab_size": 248320,
+}
+
+
+def _field(value: Any, name: str) -> Any:
+    if isinstance(value, dict):
+        return value.get(name)
+    return getattr(value, name, None)
+
+
+def assert_certified_model_config(value: Any, label: str = "base-model config") -> None:
+    model_type = _field(value, "model_type")
+    architectures = _field(value, "architectures")
+    if model_type == "qwen3_5":
+        text_config = _field(value, "text_config")
+        if (
+            not isinstance(architectures, (list, tuple))
+            or "Qwen3_5ForConditionalGeneration" not in architectures
+            or text_config is None
+        ):
+            raise ValueError(f"{label} is not the certified {CERTIFIED_BASE_MODEL} architecture")
+    elif model_type == "qwen3_5_text":
+        # AutoModelForCausalLM exposes the selected text sub-config after it
+        # dispatches the verified repository-level Qwen3.5 config.
+        text_config = value
+    else:
+        raise ValueError(f"{label} is not the certified {CERTIFIED_BASE_MODEL} architecture")
+    for name, expected in CERTIFIED_TEXT_CONFIG.items():
+        actual = _field(text_config, name)
+        if actual != expected:
+            raise ValueError(
+                f"{label} is not the certified {CERTIFIED_BASE_MODEL} architecture: "
+                f"text_config.{name} must be {expected!r}, got {actual!r}"
+            )

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

@@ -0,0 +1,192 @@
+import argparse
+import hashlib
+import json
+import os
+import re
+from pathlib import Path
+from typing import Any
+
+from model_contract import (
+    CERTIFIED_BASE_MODEL,
+    assert_certified_model_config,
+)
+
+ALLOW_PATTERNS = [
+    "*.json",
+    "*.model",
+    "*.safetensors",
+    "*.safetensors.index.json",
+    "*.tiktoken",
+    "*.txt",
+    "chat_template.jinja",
+    "merges.txt",
+    "pytorch_model*.bin",
+    "sentencepiece.bpe.model",
+    "spiece.model",
+    "tokenizer.*",
+    "vocab.*",
+]
+
+IGNORE_PATTERNS = [
+    "*.ckpt",
+    "*.gguf",
+    "*.h5",
+    "*.msgpack",
+    "*.onnx",
+    "*.ot",
+    "optimizer.pt",
+    "rng_state*.pth",
+    "scheduler.pt",
+    "trainer_state.json",
+    "training_args.bin",
+]
+
+
+def write_json(path: str, value: dict[str, Any]) -> None:
+    destination = Path(path)
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    destination.write_text(json.dumps(value, indent=2) + "\n", encoding="utf-8")
+
+
+def configure_hugging_face_cache(cache_home: str | None) -> None:
+    """Treat model_cache as HF_HOME and keep legacy overrides consistent."""
+    if not cache_home:
+        return
+    # Keep the lexical absolute path supplied by the Node runner. Path.resolve
+    # canonicalizes macOS /var to /private/var, making metadata disagree with
+    # training even though both locations identify the same directory.
+    home = Path(cache_home).expanduser().absolute()
+    hub = home / "hub"
+    os.environ.update({
+        "HF_HOME": str(home),
+        "HF_HUB_CACHE": str(hub),
+        "HUGGINGFACE_HUB_CACHE": str(hub),
+    })
+    for deprecated in (
+        "TRANSFORMERS_CACHE",
+        "PYTORCH_TRANSFORMERS_CACHE",
+        "PYTORCH_PRETRAINED_BERT_CACHE",
+    ):
+        os.environ.pop(deprecated, None)
+
+
+def hash_file(path: Path, algorithm: str, prefix: bytes = b"") -> str:
+    digest = hashlib.new(algorithm)
+    digest.update(prefix)
+    with path.open("rb") as source:
+        for chunk in iter(lambda: source.read(8 * 1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def verify_snapshot(snapshot: Path) -> tuple[list[Path], int]:
+    config = snapshot / "config.json"
+    if not config.is_file() or config.stat().st_size == 0:
+        raise ValueError(f"Cached snapshot is missing a non-empty config.json: {snapshot}")
+    try:
+        parsed_config = json.loads(config.read_text(encoding="utf-8"))
+        assert_certified_model_config(parsed_config, f"Cached config {config}")
+    except (OSError, json.JSONDecodeError) as exc:
+        raise ValueError(f"Cached snapshot has an invalid config.json: {snapshot}") from exc
+
+    weight_files = sorted(snapshot.glob("*.safetensors")) + sorted(snapshot.glob("pytorch_model*.bin"))
+    index_files = sorted(snapshot.glob("*.safetensors.index.json")) + sorted(snapshot.glob("pytorch_model*.bin.index.json"))
+    for index in index_files:
+        try:
+            weight_map = json.loads(index.read_text(encoding="utf-8")).get("weight_map", {})
+        except (OSError, json.JSONDecodeError) as exc:
+            raise ValueError(f"Cached snapshot has an invalid weight index: {index}") from exc
+        if not isinstance(weight_map, dict) or not weight_map:
+            raise ValueError(f"Cached snapshot has an empty weight index: {index}")
+        for name in set(weight_map.values()):
+            if not isinstance(name, str) or not (snapshot / name).is_file():
+                raise ValueError(f"Cached snapshot is missing an indexed weight shard: {name}")
+    if not weight_files:
+        raise ValueError(f"Cached snapshot contains no Transformers model weights: {snapshot}")
+    for weight in weight_files:
+        if weight.stat().st_size == 0:
+            raise ValueError(f"Cached snapshot contains an empty model weight: {weight}")
+
+    tokenizer_config = snapshot / "tokenizer_config.json"
+    vocabulary = [
+        snapshot / "tokenizer.json",
+        snapshot / "tokenizer.model",
+        snapshot / "sentencepiece.bpe.model",
+        snapshot / "spiece.model",
+        snapshot / "vocab.json",
+        snapshot / "tokenizer.tiktoken",
+    ]
+    if not tokenizer_config.is_file() or not any(path.is_file() and path.stat().st_size > 0 for path in vocabulary):
+        raise ValueError(f"Cached snapshot is missing tokenizer metadata or vocabulary files: {snapshot}")
+
+    files = [path for path in snapshot.rglob("*") if path.is_file()]
+    verified_blobs = 0
+    for path in files:
+        blob = path.resolve()
+        expected = blob.name.lower()
+        if len(expected) == 64 and all(character in "0123456789abcdef" for character in expected):
+            actual = hash_file(blob, "sha256")
+        elif len(expected) == 40 and all(character in "0123456789abcdef" for character in expected):
+            actual = hash_file(blob, "sha1", f"blob {blob.stat().st_size}\0".encode())
+        else:
+            continue
+        if actual != expected:
+            raise ValueError(f"Cached Hugging Face blob checksum mismatch: {path}")
+        verified_blobs += 1
+    return files, verified_blobs
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Prefetch a TT Local Hugging Face base model.")
+    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"))
+    base_model = str(payload["base_model"])
+    if base_model != CERTIFIED_BASE_MODEL:
+        raise ValueError(
+            f"The bundled trainer currently certifies only {CERTIFIED_BASE_MODEL}; got {base_model!r}"
+        )
+    cache_dir = payload.get("model_cache")
+    configure_hugging_face_cache(cache_dir)
+
+    # Import only after the cache environment is final. huggingface_hub reads
+    # these variables while importing its constants module.
+    from huggingface_hub import constants, snapshot_download
+
+    action = "Verifying cached" if payload.get("local_files_only") else "Prefetching"
+    print(f"{action} {base_model} in {constants.HF_HUB_CACHE}...", flush=True)
+    snapshot_path = snapshot_download(
+        repo_id=base_model,
+        revision=payload.get("revision"),
+        token=os.getenv("HF_TOKEN"),
+        allow_patterns=ALLOW_PATTERNS,
+        ignore_patterns=IGNORE_PATTERNS,
+        local_files_only=bool(payload.get("local_files_only", False)),
+    )
+    snapshot = Path(snapshot_path)
+    snapshot_revision = snapshot.name
+    if not re.fullmatch(r"[0-9a-fA-F]{40}", snapshot_revision):
+        raise ValueError(
+            "Hugging Face snapshot did not resolve to a 40-character immutable commit SHA"
+        )
+    print(f"Verifying snapshot files under {snapshot}...", flush=True)
+    snapshot_files, verified_blob_count = verify_snapshot(snapshot)
+
+    write_json(args.output, {
+        "ok": True,
+        "base_model": base_model,
+        "model_cache": str(constants.HF_HOME),
+        "hf_home": str(constants.HF_HOME),
+        "hub_cache": str(constants.HF_HUB_CACHE),
+        "snapshot_path": snapshot_path,
+        "snapshot_revision": snapshot_revision.lower(),
+        "file_count": len(snapshot_files),
+        "size_bytes": sum(path.stat().st_size for path in snapshot_files),
+        "verified_blob_count": verified_blob_count,
+    })
+
+
+if __name__ == "__main__":
+    main()