From 5afbe9c8555546e09ccb54bf2b5612d652b8dc6e Mon Sep 17 00:00:00 2001 From: William Gao Date: Tue, 30 Jun 2026 08:35:13 -0700 Subject: [PATCH 1/2] feat(smoke-test): mount qwen3-0.6b base model via weights API instead of HF The smoke test pulled Qwen/Qwen3-0.6B directly from HuggingFace inside the training container on every run, which is getting rate-limited by HF. Mount the base model through the truss-train weights API instead: Baseten mirrors it into the Baseten Delivery Network once and serves it from cache on subsequent runs, so repeated CI runs no longer hit HF. - config.py: add WeightsSource(hf://Qwen/Qwen3-0.6B) mounted at /mnt/user/Qwen3-0.6B, and pass the mount path to the container via the BASE_MODEL_PATH env var so it stays defined in one place. - train.py: read BASE_MODEL_PATH and point axolotl base_model at the mounted path. --- .../qwen3-0.6b-axolotl/training/config.py | 15 ++++++++-- examples/qwen3-0.6b-axolotl/training/train.py | 28 ++++++++++++++++++- 2 files changed, 39 insertions(+), 4 deletions(-) diff --git a/examples/qwen3-0.6b-axolotl/training/config.py b/examples/qwen3-0.6b-axolotl/training/config.py index 584bbf9..3fdfde5 100644 --- a/examples/qwen3-0.6b-axolotl/training/config.py +++ b/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, + ) + ], ) diff --git a/examples/qwen3-0.6b-axolotl/training/train.py b/examples/qwen3-0.6b-axolotl/training/train.py index 5daa133..185972d 100644 --- a/examples/qwen3-0.6b-axolotl/training/train.py +++ b/examples/qwen3-0.6b-axolotl/training/train.py @@ -5,6 +5,8 @@ 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 @@ -12,10 +14,34 @@ 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", From 1c682cc1c912004166c77ea4c4907fb5ad6f0a51 Mon Sep 17 00:00:00 2001 From: William Gao Date: Tue, 30 Jun 2026 08:50:03 -0700 Subject: [PATCH 2/2] fix(smoke-test): pass run_id to create_model_version_from_inference_template truss's create_model_version_from_inference_template() gained a required `run_id` positional argument, so the deploy-and-infer step of the smoke test fails with "missing 1 required positional argument: 'run_id'" even when the training job completes and checkpoints successfully. The smoke test deploys training-job checkpoints, and _hydrate_deploy_config raises if run_id is set alongside training-job checkpoints, so the correct value here is None. --- bin/test_example.py | 1 + 1 file changed, 1 insertion(+) diff --git a/bin/test_example.py b/bin/test_example.py index e34e6c7..8d0ce9b 100755 --- a/bin/test_example.py +++ b/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, )