@tuned-tensor/cli 0.11.0 → 0.12.0-beta.0

package/docs/local-runtime/spark.md

@@ -4,8 +4,10 @@
 > `tt-local` is deprecated. Install `@tuned-tensor/cli` and use `tt local`.
 > Existing projects still work.
 
-DGX Spark is the reference host for the local workflow's first supported path:
-`Qwen/Qwen3.5-2B` text SFT with a LoRA adapter.
+DGX Spark is the reference host for local text SFT with LoRA adapters. TT Local
+certifies `Qwen/Qwen3.5-2B`, the larger
+`nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16` checkpoint, and
+`meta-models/Muse-Glimmer-30B`.
 
 ## Check the host
 
@@ -28,6 +30,41 @@ cd ~/tuned-tensor-runs/support-adapter
 tt local init --name "Support Adapter" --model Qwen/Qwen3.5-2B --profile spark
 ```
 
+For a Nemotron execution worker:
+
+```bash
+mkdir -p ~/tuned-tensor-runs/nemotron-worker
+cd ~/tuned-tensor-runs/nemotron-worker
+tt local init \
+  --name "Nemotron Worker" \
+  --model nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 \
+  --profile spark
+```
+
+Nemotron uses the BF16 customization checkpoint rather than the NVFP4
+deployment checkpoint. Its defaults enable activation checkpointing, cap the
+sequence length at 1,024 tokens, and apply LoRA only to shared attention and
+Mamba projections. The base weights require about 66 GB before cache and
+training overhead, so this path is not intended for ordinary consumer GPUs.
+
+For a Muse Glimmer execution worker:
+
+```bash
+mkdir -p ~/tuned-tensor-runs/muse-glimmer-worker
+cd ~/tuned-tensor-runs/muse-glimmer-worker
+tt local init \
+  --name "Muse Glimmer Worker" \
+  --model meta-models/Muse-Glimmer-30B \
+  --profile spark
+```
+
+Muse Glimmer is a vision-language checkpoint. The runtime fine-tunes its text
+tower only: it loads the checkpoint through `AutoModelForImageTextToText` (the
+text tower is not exposed in the causal-LM auto mapping) and leaves the vision
+tower frozen and unused. Defaults enable activation checkpointing and cap the
+sequence length at 2,048 tokens; the base weights are comparable to Nemotron
+and require Spark-class unified memory.
+
 Edit both generated examples in `tunedtensor.json`. For a meaningful run,
 replace them with a larger, representative dataset and a separate validation
 split.

package/package.json

@@ -1,7 +1,11 @@
 {
   "name": "@tuned-tensor/cli",
-  "version": "0.11.0",
+  "version": "0.12.0-beta.0",
   "description": "Local Pi agent and unified cloud/local CLI for Tuned Tensor workflows",
+  "repository": {
+    "type": "git",
+    "url": "git+https://github.com/tunedtensor/tuned-tensor-cli.git"
+  },
   "type": "module",
   "bin": {
     "tt": "dist/index.js"

package/training/local-runner/pyproject.toml

@@ -5,10 +5,10 @@ requires-python = ">=3.11"
 dependencies = [
   "accelerate>=1.0.0",
   "huggingface-hub>=0.36.0",
-  "peft>=0.13.0",
+  "peft>=0.20.0",
   "safetensors>=0.4.5",
   "torch>=2.11.0",
-  "transformers>=4.57.0",
+  "transformers>=5.15.0",
 ]
 
 [tool.uv.sources]

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

@@ -10,8 +10,12 @@ from tempfile import TemporaryDirectory
 from typing import Any
 
 from model_contract import (
+    CERTIFIED_BASE_MODELS,
     CERTIFIED_BASE_MODEL,
+    assert_certified_base_model,
+    assert_certified_base_model_revision,
     assert_certified_model_config,
+    chat_template_kwargs,
 )
 
 MAX_ARCHIVE_MEMBERS = 20_000
@@ -46,18 +50,20 @@ def configure_hugging_face_cache(cache_home: str | None) -> None:
 
 def import_runtime_dependencies() -> None:
     """Import libraries only after Hugging Face cache settings are final."""
-    global torch, PeftModel, AutoModelForCausalLM, AutoTokenizer
+    global torch, PeftModel, AutoModelForCausalLM, AutoModelForImageTextToText, AutoTokenizer
 
     import torch as torch_module
     from peft import PeftModel as peft_model
     from transformers import (
         AutoModelForCausalLM as auto_model_for_causal_lm,
+        AutoModelForImageTextToText as auto_model_for_image_text_to_text,
         AutoTokenizer as auto_tokenizer,
     )
 
     torch = torch_module
     PeftModel = peft_model
     AutoModelForCausalLM = auto_model_for_causal_lm
+    AutoModelForImageTextToText = auto_model_for_image_text_to_text
     AutoTokenizer = auto_tokenizer
 
 
@@ -111,6 +117,7 @@ def format_prompt(
     tokenizer: Any,
     system: str,
     prompt: str,
+    base_model: str = CERTIFIED_BASE_MODEL,
 ) -> str:
     messages = [
         {"role": "system", "content": system},
@@ -121,9 +128,10 @@ def format_prompt(
             messages,
             tokenize=False,
             add_generation_prompt=True,
+            **chat_template_kwargs(base_model),
         )
     except Exception as exc:
-        raise ValueError(f"Qwen chat-template rendering failed: {exc}") from exc
+        raise ValueError(f"Model chat-template rendering failed: {exc}") from exc
 
 
 def resolve_device(device: str) -> str:
@@ -135,30 +143,37 @@ def resolve_device(device: str) -> str:
 
 
 def _assert_certified_model_source(base_model: str) -> None:
-    if not Path(base_model).exists() and base_model != CERTIFIED_BASE_MODEL:
+    if not Path(base_model).exists() and base_model not in CERTIFIED_BASE_MODELS:
         raise ValueError(
-            f"The bundled evaluator currently certifies only {CERTIFIED_BASE_MODEL}; got {base_model!r}"
+            f"The bundled evaluator does not certify {base_model!r}"
         )
 
 
-def load_text_model(payload: dict[str, Any], adapter_path: str | None):
+def load_text_model(
+    payload: dict[str, Any],
+    adapter_path: str | None,
+    model_loader: str = "causal_lm",
+):
     base_model = str(payload["base_model"])
-    _assert_certified_model_source(base_model)
+    model_source = str(payload.get("model_source", base_model))
+    assert_certified_base_model(base_model, "Evaluation base model")
+    _assert_certified_model_source(model_source)
+    revision = payload.get("base_model_revision")
+    assert_certified_base_model_revision(base_model, revision, "Evaluation base model revision")
     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():
+    if not Path(model_source).exists():
         source_kwargs["local_files_only"] = True
-    if revision and not Path(base_model).exists():
+    if revision and not Path(model_source).exists():
         source_kwargs["revision"] = revision
 
-    tokenizer = AutoTokenizer.from_pretrained(base_model, **source_kwargs)
+    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")
+        raise ValueError(f"{base_model} tokenizer has no EOS token")
     if tokenizer.pad_token_id is None:
         tokenizer.pad_token = tokenizer.eos_token
 
@@ -172,8 +187,13 @@ def load_text_model(payload: dict[str, Any], adapter_path: str | None):
     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)
+    model_class = (
+        AutoModelForImageTextToText
+        if model_loader == "image_text_to_text"
+        else AutoModelForCausalLM
+    )
+    model = model_class.from_pretrained(model_source, **model_kwargs)
+    assert_certified_model_config(model.config, expected_model_id=base_model)
     if adapter_path:
         model = PeftModel.from_pretrained(model, adapter_path)
     if device != "cuda":
@@ -199,11 +219,12 @@ def generate_text_one(
     system: str,
     example: dict[str, Any],
     generation: dict[str, Any],
+    base_model: str = CERTIFIED_BASE_MODEL,
 ) -> dict[str, Any]:
     if example.get("input_assets"):
-        raise ValueError("The bundled Qwen SFT evaluator is text-only and does not accept input_assets")
+        raise ValueError("The bundled SFT evaluator is text-only and does not accept input_assets")
     prompt = str(example["input"])
-    formatted = format_prompt(tokenizer, system, prompt)
+    formatted = format_prompt(tokenizer, system, prompt, base_model)
     inputs = tokenizer(formatted, return_tensors="pt")
     target_device = next(model.parameters()).device
     inputs = {key: value.to(target_device) for key, value in inputs.items()}
@@ -235,8 +256,12 @@ def main() -> None:
     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")
+    model_loader = str(payload.get("model_loader") or "causal_lm")
+    if model_loader not in ("causal_lm", "image_text_to_text"):
+        raise ValueError(
+            f"The bundled evaluator does not support model_loader={model_loader!r}; "
+            "expected causal_lm or image_text_to_text"
+        )
     configure_hugging_face_cache(payload.get("model_cache"))
     import_runtime_dependencies()
 
@@ -245,7 +270,7 @@ def main() -> None:
         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)
+        model, tokenizer, _device = load_text_model(payload, adapter_path, model_loader)
         results = [
             generate_text_one(
                 model,
@@ -253,6 +278,7 @@ def main() -> None:
                 str(payload.get("system", "")),
                 example,
                 generation,
+                str(payload["base_model"]),
             )
             for example in payload.get("examples", [])
         ]

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

@@ -4,7 +4,16 @@ from typing import Any
 
 
 CERTIFIED_BASE_MODEL = "Qwen/Qwen3.5-2B"
-CERTIFIED_TEXT_CONFIG = {
+NEMOTRON_BASE_MODEL = "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16"
+MUSE_GLIMMER_BASE_MODEL = "meta-models/Muse-Glimmer-30B"
+# Immutable Hugging Face revision reviewed for Nemotron local fine-tuning.
+# Node validates local snapshots against this revision and its certified file
+# digests before invoking Python.
+NEMOTRON_BASE_MODEL_REVISION = "ce38b6ab8b252b4b8ee7165b4605e93191cafd73"
+# Immutable Hugging Face revision reviewed and certified for Muse Glimmer.
+MUSE_GLIMMER_BASE_MODEL_REVISION = "a4e59da52a7bc87ae7251dd5545c0dd437c44b68"
+CERTIFIED_BASE_MODELS = (CERTIFIED_BASE_MODEL, NEMOTRON_BASE_MODEL, MUSE_GLIMMER_BASE_MODEL)
+CERTIFIED_QWEN_TEXT_CONFIG = {
     "model_type": "qwen3_5_text",
     "hidden_size": 2048,
     "num_hidden_layers": 24,
@@ -13,6 +22,27 @@ CERTIFIED_TEXT_CONFIG = {
     "intermediate_size": 6144,
     "vocab_size": 248320,
 }
+CERTIFIED_MUSE_GLIMMER_TEXT_CONFIG = {
+    "model_type": "muse_glimmer_text",
+    "hidden_size": 6656,
+    "num_hidden_layers": 52,
+    "num_attention_heads": 32,
+    "num_key_value_heads": 2,
+    "intermediate_size": 19968,
+    "vocab_size": 202048,
+}
+CERTIFIED_NEMOTRON_CONFIG = {
+    "hidden_size": 2688,
+    "num_hidden_layers": 52,
+    "num_attention_heads": 32,
+    "num_key_value_heads": 2,
+    "intermediate_size": 1856,
+    "vocab_size": 131072,
+    "n_routed_experts": 128,
+    "num_experts_per_tok": 6,
+    "num_nextn_predict_layers": 1,
+    "max_position_embeddings": 262144,
+}
 
 
 def _field(value: Any, name: str) -> Any:
@@ -21,27 +51,117 @@ def _field(value: Any, name: str) -> Any:
     return getattr(value, name, None)
 
 
-def assert_certified_model_config(value: Any, label: str = "base-model config") -> None:
+def assert_certified_model_config(
+    value: Any,
+    label: str = "base-model config",
+    expected_model_id: str | None = None,
+) -> None:
     model_type = _field(value, "model_type")
     architectures = _field(value, "architectures")
-    if model_type == "qwen3_5":
-        text_config = _field(value, "text_config")
+
+    if model_type in ("qwen3_5", "qwen3_5_text"):
+        actual_model_id = CERTIFIED_BASE_MODEL
+        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")
+        else:
+            # AutoModelForCausalLM exposes the selected text sub-config after it
+            # dispatches the verified repository-level Qwen3.5 config.
+            text_config = value
+        certified = CERTIFIED_QWEN_TEXT_CONFIG
+    elif model_type in ("muse_glimmer", "muse_glimmer_text"):
+        actual_model_id = MUSE_GLIMMER_BASE_MODEL
+        if model_type == "muse_glimmer":
+            text_config = _field(value, "text_config")
+            if (
+                not isinstance(architectures, (list, tuple))
+                or "MuseGlimmerForConditionalGeneration" not in architectures
+                or text_config is None
+            ):
+                raise ValueError(f"{label} is not the certified {MUSE_GLIMMER_BASE_MODEL} architecture")
+        else:
+            text_config = value
+        certified = CERTIFIED_MUSE_GLIMMER_TEXT_CONFIG
+    elif model_type == "nemotron_h":
+        actual_model_id = NEMOTRON_BASE_MODEL
         if (
             not isinstance(architectures, (list, tuple))
-            or "Qwen3_5ForConditionalGeneration" not in architectures
-            or text_config is None
+            or "NemotronHForCausalLM" not in architectures
         ):
-            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
+            raise ValueError(f"{label} is not the certified {NEMOTRON_BASE_MODEL} architecture")
+        for name, expected in CERTIFIED_NEMOTRON_CONFIG.items():
+            actual = _field(value, name)
+            if actual != expected:
+                raise ValueError(
+                    f"{label} is not the certified {NEMOTRON_BASE_MODEL} architecture: "
+                    f"{name} must be {expected!r}, got {actual!r}"
+                )
+        if expected_model_id and expected_model_id != actual_model_id:
+            raise ValueError(
+                f"{label} does not match requested base model {expected_model_id}; "
+                f"it matches {actual_model_id}"
+            )
+        return
     else:
-        raise ValueError(f"{label} is not the certified {CERTIFIED_BASE_MODEL} architecture")
-    for name, expected in CERTIFIED_TEXT_CONFIG.items():
+        raise ValueError(f"{label} is not a certified TT Local base-model architecture")
+
+    for name, expected in certified.items():
         actual = _field(text_config, name)
         if actual != expected:
             raise ValueError(
-                f"{label} is not the certified {CERTIFIED_BASE_MODEL} architecture: "
+                f"{label} is not the certified {actual_model_id} architecture: "
                 f"text_config.{name} must be {expected!r}, got {actual!r}"
             )
+    if expected_model_id and expected_model_id != actual_model_id:
+        raise ValueError(
+            f"{label} does not match requested base model {expected_model_id}; "
+            f"it matches {actual_model_id}"
+        )
+
+
+def assert_certified_base_model(model_id: str, label: str = "base model") -> None:
+    if model_id not in CERTIFIED_BASE_MODELS:
+        supported = ", ".join(CERTIFIED_BASE_MODELS)
+        raise ValueError(f"{label} must be one of {supported}; got {model_id!r}")
+
+
+def assert_certified_base_model_revision(base_model: str, revision: str | None, label: str = "base model revision") -> None:
+    """Require the certified immutable revision for pinned training models.
+
+    Qwen remains unpinned for backward compatibility. Nemotron and Muse
+    Glimmer loads are bound here; local snapshot contents are additionally
+    verified by the Node runtime.
+    """
+    expected = {
+        NEMOTRON_BASE_MODEL: NEMOTRON_BASE_MODEL_REVISION,
+        MUSE_GLIMMER_BASE_MODEL: MUSE_GLIMMER_BASE_MODEL_REVISION,
+    }.get(base_model)
+    if expected is None:
+        return
+    if revision is None:
+        raise ValueError(f"{label} is required for {base_model}")
+    if revision != expected:
+        raise ValueError(
+            f"{label} must be {expected} for {base_model}; got {revision!r}"
+        )
+
+
+def parse_lora_target_modules(value: str | None) -> str | list[str]:
+    raw = (value or "all-linear").strip()
+    if raw == "all-linear":
+        return raw
+    modules = [item.strip() for item in raw.split(",")]
+    if not modules or any(not item for item in modules):
+        raise ValueError("LoRA target modules must be 'all-linear' or a comma-separated non-empty list")
+    return modules
+
+
+def chat_template_kwargs(model_id: str) -> dict[str, Any]:
+    if model_id == NEMOTRON_BASE_MODEL:
+        return {"enable_thinking": False}
+    return {}

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

@@ -7,11 +7,17 @@ from pathlib import Path
 from typing import Any
 
 from model_contract import (
-    CERTIFIED_BASE_MODEL,
+    assert_certified_base_model,
+    assert_certified_base_model_revision,
     assert_certified_model_config,
 )
 
 ALLOW_PATTERNS = [
+    ".gitattributes",
+    "LICENSE",
+    "README.md",
+    "*.md",
+    "*.png",
     "*.json",
     "*.model",
     "*.safetensors",
@@ -79,13 +85,17 @@ def hash_file(path: Path, algorithm: str, prefix: bytes = b"") -> str:
     return digest.hexdigest()
 
 
-def verify_snapshot(snapshot: Path) -> tuple[list[Path], int]:
+def verify_snapshot(snapshot: Path, expected_model_id: str | None = None) -> 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}")
+        assert_certified_model_config(
+            parsed_config,
+            f"Cached config {config}",
+            expected_model_id=expected_model_id,
+        )
     except (OSError, json.JSONDecodeError) as exc:
         raise ValueError(f"Cached snapshot has an invalid config.json: {snapshot}") from exc
 
@@ -144,10 +154,13 @@ def main() -> None:
 
     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}"
-        )
+    assert_certified_base_model(base_model, "Prefetch base model")
+    requested_revision = payload.get("revision")
+    assert_certified_base_model_revision(
+        base_model,
+        requested_revision,
+        "Prefetch base model revision",
+    )
     cache_dir = payload.get("model_cache")
     configure_hugging_face_cache(cache_dir)
 
@@ -171,8 +184,12 @@ def main() -> None:
         raise ValueError(
             "Hugging Face snapshot did not resolve to a 40-character immutable commit SHA"
         )
+    if requested_revision and snapshot_revision.lower() != str(requested_revision).lower():
+        raise ValueError(
+            f"Hugging Face returned revision {snapshot_revision}, not requested revision {requested_revision}"
+        )
     print(f"Verifying snapshot files under {snapshot}...", flush=True)
-    snapshot_files, verified_blob_count = verify_snapshot(snapshot)
+    snapshot_files, verified_blob_count = verify_snapshot(snapshot, base_model)
 
     write_json(args.output, {
         "ok": True,

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

@@ -20,10 +20,12 @@ from evaluate import (
     resolve_adapter_path,
     sampling_kwargs,
 )
+from model_contract import assert_certified_base_model_revision, chat_template_kwargs
 
 
 MODEL_ARTIFACT = os.environ.get("TT_MODEL_ARTIFACT")
 BASE_MODEL = os.environ["TT_BASE_MODEL"]
+MODEL_SOURCE = os.environ.get("TT_MODEL_SOURCE", 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")
@@ -42,8 +44,12 @@ 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")
+if MODEL_LOADER not in ("causal_lm", "image_text_to_text"):
+    raise ValueError(
+        f"The bundled model server does not support TT_MODEL_LOADER={MODEL_LOADER!r}; "
+        "expected causal_lm or image_text_to_text"
+    )
+assert_certified_base_model_revision(BASE_MODEL, BASE_MODEL_REVISION, "Serving base model revision")
 configure_hugging_face_cache(os.environ.get("HF_HOME"))
 import_runtime_dependencies()
 TEMP_DIR = TemporaryDirectory(prefix="tt-local-serve-")
@@ -54,11 +60,12 @@ ADAPTER_PATH = (
 )
 MODEL_PAYLOAD = {
     "base_model": BASE_MODEL,
+    "model_source": MODEL_SOURCE,
     "base_model_revision": BASE_MODEL_REVISION,
     "device": DEVICE_REQUEST,
     "model_loader": MODEL_LOADER,
 }
-MODEL, TOKENIZER, DEVICE = load_text_model(MODEL_PAYLOAD, ADAPTER_PATH)
+MODEL, TOKENIZER, DEVICE = load_text_model(MODEL_PAYLOAD, ADAPTER_PATH, MODEL_LOADER)
 GENERATION_LOCK = threading.Lock()
 REQUEST_SLOTS = threading.BoundedSemaphore(MAX_CONCURRENT_REQUESTS)
 
@@ -73,7 +80,7 @@ def text_content(content: Any) -> str:
     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.")
+            raise ValueError("The bundled 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.")
@@ -104,7 +111,7 @@ def normalize_messages(raw_messages: Any) -> list[dict[str, str]]:
         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
+        # The certified chat templates accept 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)})
@@ -120,9 +127,10 @@ def generate_text(messages: list[dict[str, str]], generation: dict[str, Any]) ->
             messages,
             tokenize=False,
             add_generation_prompt=True,
+            **chat_template_kwargs(BASE_MODEL),
         )
     except Exception as exc:
-        raise ValueError(f"Qwen chat-template rendering failed: {exc}") from exc
+        raise ValueError(f"Model 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.")

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

@@ -23,6 +23,7 @@ def build_assistant_only_example(
     messages: Any,
     *,
     max_length: int,
+    chat_template_kwargs: Mapping[str, Any] | None = None,
 ) -> dict[str, list[int]]:
     """Tokenize one chat row while computing loss only on the final assistant turn.
 
@@ -52,12 +53,14 @@ def build_assistant_only_example(
         raise ValueError("the final assistant message must contain non-empty text")
 
     try:
+        template_kwargs = dict(chat_template_kwargs or {})
         prompt_ids = _token_ids(
             tokenizer.apply_chat_template(
                 messages[:-1],
                 tokenize=True,
                 add_generation_prompt=True,
                 return_dict=False,
+                **template_kwargs,
             ),
             "prompt chat template",
         )
@@ -67,19 +70,20 @@ def build_assistant_only_example(
                 tokenize=True,
                 add_generation_prompt=False,
                 return_dict=False,
+                **template_kwargs,
             ),
             "full chat template",
         )
     except Exception as exc:
-        raise ValueError(f"Qwen chat-template rendering failed: {exc}") from exc
+        raise ValueError(f"Model 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"
+            "Model 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")
+        raise ValueError("Model 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}; "