feat(smoke-test): mount qwen3-0.6b base model via weights API instead of HF

bin/test_example.py

@@ -488,6 +488,7 @@ def deploy_and_infer(
         deploy_config,
         project_id=project_id,
         job_id=job_id,
+        run_id=None,
         dry_run=False,
     )
 

examples/qwen3-0.6b-axolotl/training/config.py

@@ -1,9 +1,11 @@
 # Import necessary classes from the Baseten Training SDK
-from truss_train import definitions
+from truss_train import definitions, WeightsSource
 from truss.base import truss_config
 
 project_name = "demo/qwen3-0.6b"
 
+MODEL_MOUNT_PATH = "/mnt/user/Qwen3-0.6B"
+
 # 1. Define a base image for your training job
 BASE_IMAGE = "axolotlai/axolotl:main-20250811-py3.11-cu126-2.7.1"
 
@@ -19,8 +21,9 @@
         f"axolotl fetch deepspeed_configs && torchrun --nproc-per-node={NUM_GPUS} train.py",
     ],
     environment_variables={
-        # Secrets (ensure these are configured in your Baseten workspace)
-        # Include other environment variables as needed
+        # train.py seeds the HF cache from this mount so the model loads from
+        # the mirror instead of HuggingFace, while keeping its canonical repo id.
+        "MODEL_MOUNT_PATH": MODEL_MOUNT_PATH,
     },
     cache_config=definitions.CacheConfig(
         enabled=True,
@@ -45,6 +48,12 @@
     image=definitions.Image(base_image=BASE_IMAGE),
     compute=training_compute,
     runtime=training_runtime,
+    weights=[
+        WeightsSource(
+            source="hf://Qwen/Qwen3-0.6B",
+            mount_location=MODEL_MOUNT_PATH,
+        )
+    ],
 )
 
 

examples/qwen3-0.6b-axolotl/training/train.py

@@ -5,17 +5,43 @@
 if CACHE_DIR:
     os.environ["HF_HOME"] = CACHE_DIR
 
+# Imported after HF_HOME is set so the hub cache constants resolve to it.
+from huggingface_hub import HfApi
 from axolotl.utils.dict import DictDefault
 from axolotl.cli.config import load_cfg
 from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 
 OUTPUT_DIR = os.environ.get("BT_CHECKPOINT_DIR", "outputs/qwen3-0.6b")
 
+MODEL_MOUNT_PATH = os.environ["MODEL_MOUNT_PATH"]
+MODEL_ID = "Qwen/Qwen3-0.6B"
+
+
+def seed_hf_cache_from_mount(model_id: str, mount_path: str) -> None:
+    """Expose the BDN-mounted weights to HF under the canonical repo id.
+
+    Lets from_pretrained(model_id) load the weights from the mount instead of
+    downloading them, while still recording the HF repo id in the saved LoRA
+    adapter — a deployable checkpoint needs base_model to be 'namespace/model'
+    (resolved from the BDN mirror), not a local path. Recording it at train
+    time avoids the post-hoc re-sync/re-discovery race a later rewrite hits."""
+    hub = Path(os.environ.get("HF_HOME", Path.home() / ".cache" / "huggingface")) / "hub"
+    repo_dir = hub / f"models--{model_id.replace('/', '--')}"
+    commit = HfApi().model_info(model_id).sha
+    snapshot = repo_dir / "snapshots" / commit
+    (repo_dir / "refs").mkdir(parents=True, exist_ok=True)
+    (repo_dir / "refs" / "main").write_text(commit)
+    snapshot.parent.mkdir(parents=True, exist_ok=True)
+    if not snapshot.exists():
+        snapshot.symlink_to(mount_path)
+
+
 def main():
+    seed_hf_cache_from_mount(MODEL_ID, MODEL_MOUNT_PATH)
     config = DictDefault(
         adapter="qlora",
-        base_model="Qwen/Qwen3-0.6B",
+        base_model=MODEL_ID,
         bf16=True,
         # chat_template="tokenizer_default_fallback_chatml",