@tuned-tensor/cli 0.24.0 → 0.26.0

package/training/adapter/src/prefetch.py

@@ -137,6 +137,16 @@ def verify_snapshot(snapshot: Path, expected_model_id: str | None = None) -> tup
     return files, verified_blobs
 
 
+def check_requested_revision(base_model: str, revision: str | None, resolve_latest: bool) -> None:
+    """Certified models keep their reviewed revision. An uncertified model needs
+    a pinned commit, except for a direct download that resolves its latest
+    commit (`tt models prefetch --repo`); the resolved SHA is reported so the
+    spec can pin it."""
+    if resolve_latest and revision is None and not is_certified_base_model(base_model):
+        return
+    assert_certified_base_model_revision(base_model, revision, "Prefetch base model revision")
+
+
 def main() -> None:
     parser = argparse.ArgumentParser(description="Prefetch a TT Hugging Face base model for local use.")
     parser.add_argument("--input", required=True)
@@ -147,11 +157,7 @@ def main() -> None:
     base_model = str(payload["base_model"])
     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",
-    )
+    check_requested_revision(base_model, requested_revision, bool(payload.get("resolve_latest")))
     cache_dir = payload.get("model_cache")
     configure_hugging_face_cache(cache_dir)
 

package/training/adapter/src/serve.py

@@ -63,6 +63,35 @@ def pin_native_content_format(args: list[str], model_source: str) -> None:
     args += ["--chat-template-content-format", content_format]
 
 
+def prepare_serving_adapter(adapter_path: str, temp_dir: Path) -> str:
+    # Transformers trains Qwen's text-only model, but vLLM loads the repository's
+    # conditional architecture, even with --language-model-only. Its LoRA mapper
+    # expects the language_model prefix; otherwise it silently misses our layers.
+    if os.environ.get("TT_BASE_MODEL") != "Qwen/Qwen3.5-2B":
+        return adapter_path
+    from safetensors import safe_open
+    from safetensors.torch import save_file
+
+    source = Path(adapter_path)
+    old = "base_model.model.model.layers."
+    new = "base_model.model.model.language_model.layers."
+    with safe_open(str(source / "adapter_model.safetensors"), framework="pt", device="cpu") as weights:
+        if not any(key.startswith(old) for key in weights.keys()):
+            return adapter_path
+        tensors = {}
+        for key in weights.keys():
+            mapped = new + key[len(old):] if key.startswith(old) else key
+            if mapped in tensors:
+                raise ValueError("Adapter contains conflicting text and conditional layer names.")
+            tensors[mapped] = weights.get_tensor(key)
+        destination = temp_dir / "serving-adapter"
+        destination.mkdir()
+        save_file(tensors, str(destination / "adapter_model.safetensors"), metadata=weights.metadata())
+    # Keep the verified archive and extracted training layout unchanged.
+    (destination / "adapter_config.json").write_bytes((source / "adapter_config.json").read_bytes())
+    return str(destination)
+
+
 def build_vllm_args(model_source: str, adapter_path: str | None, temp_dir: Path) -> list[str]:
     args = [
         "serve", model_source,
@@ -101,6 +130,7 @@ def build_vllm_args(model_source: str, adapter_path: str | None, temp_dir: Path)
         adapter_base = config.get("base_model_name_or_path")
         if adapter_base not in {os.environ["TT_BASE_MODEL"], model_source}:
             raise ValueError("The adapter base does not match the verified serving model.")
+        adapter_path = prepare_serving_adapter(adapter_path, temp_dir)
         base_name = "base:" + os.environ["TT_BASE_MODEL"]
         args[args.index("--served-model-name") + 1] = base_name
         args += ["--enable-lora", "--max-lora-rank", str(max(8, int(config["r"]))),

package/training/adapter/src/train.py

@@ -26,6 +26,7 @@ from model_contract import (
     assert_certified_base_model_revision,
     assert_certified_model_config,
     chat_template_kwargs,
+    cuda_model_load_device_map,
     parse_lora_target_modules,
     uncertified_load_kwargs,
 )
@@ -216,13 +217,16 @@ def create_model_and_tokenizer(model_source: str) -> tuple[Any, Any, torch.dtype
         if hp("model_loader", "causal_lm") == "image_text_to_text"
         else AutoModelForCausalLM
     )
+    load_device_map = cuda_model_load_device_map(torch)
     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()},
+        device_map=load_device_map,
     )
+    if load_device_map is None:
+        model = model.to("cuda")
     assert_certified_model_config(
         model.config,
         expected_model_id=str(hp("base_model", CERTIFIED_BASE_MODEL)),