Hugging Face

gpu-train integrates with the Hugging Face Hub in two directions: pull baseline models to a local cache, and inject your token into every run so the training script can download gated models from the box (traffic is outbound, so it works through NAT — just like W&B).

Install the extra:

pip install "gpu-train[hf]"

Pull a baseline model#

import gpu_train
 
# Download to ~/.gpu-train/models/<model_id> (or pass local_dir=...)
path = gpu_train.hf.pull("meta-llama/Llama-3.2-1B")
 
# Pin a revision / subset of files:
path = gpu_train.hf.pull(
    "meta-llama/Llama-3.2-1B",
    revision="main",
    allow_patterns=["*.safetensors", "*.json"],
)

pull(
    model_id: str,
    *,
    revision: str | None = None,
    local_dir: str | Path | None = None,
    allow_patterns: list[str] | None = None,
    token: str | None = None,
) -> Path

Already have a model on disk? Adopt it as a baseline (validated, returns the resolved path):

base = gpu_train.hf.from_local("./my-checkpoint")

Token injection into jobs#

Configure a token once and gpu-train injects HF_TOKEN / HUGGING_FACE_HUB_TOKEN (and optionally HF_HOME) into every job, so your train.py can from_pretrained(...) gated models on the remote box. Tokens are resolved (never stored in the package) and written to owner-only files on the box, like other secrets.

Configure it any of these ways:

configure(
    registry={
        "integrations": {
            "huggingface": {
                "token_ref": "env://HF_TOKEN",   # or store://hf/token
                "cache_dir": "/workspace/hf",     # optional -> HF_HOME
                "enabled": True,
            }
        },
    },
    use=[...],
)

• Set HF_TOKEN (or HUGGING_FACE_HUB_TOKEN) in the environment, or
• Connect it from the dashboard's Hugging Face card (saved locally, chmod 600).

In the dashboard#

The Providers & integrations page has a Hugging Face card to connect/disconnect a token and a pull a baseline model form. Relevant endpoints: