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, ) 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",