AutoGineer
Generated banner for Tao Finetune Huggingface Model
Generated logo for Tao Finetune Huggingface Model

Tao Finetune Huggingface Model

NVIDIAOfficialCreated by NVIDIA Corporation

  • Skill
  • Workflow
  • Software development
  • Apache-2.0
  • Medium risk
Download ZIP

From the author

Fine-tune any HuggingFace CV / VLM / LLM model on local NVIDIA GPUs inside an NGC PyTorch container when no dedicated TAO model skill matches. Use when the user wants to fine-tune a HuggingFace model (full or LoRA), train a vision / VLM / LLM model end-to-end, generate a reproducible HF training pipeline, smoke-test a HuggingFace model locally before scale-up, push a fine-tuned model to the HF Hub with a model card, or emit a self-contained rerun skill for an existing HuggingFace finetune. Supports image classification, object detection, semantic / instance / panoptic segmentation, depth estimation, image-text-to-text VLM (SFT / LoRA), and LLM SFT / DPO / GRPO. Six-step workflow: inspect and qualify, hardware and NGC image, research, generate and smoke, train + eval + infer, push and emit rerun skill. Do not use for any Hugging Face model ID claimed by a dedicated `skills/models/*` skill; the model skill and its declared execution environment take precedence.
— the description by NVIDIA Corporation

AutoGineer summary

Written by AutoGineer from the skill’s own files, not by the author

This skill provides a structured, six-step workflow for fine-tuning HuggingFace CV, VLM, and LLM models on local NVIDIA GPUs. It leverages NGC PyTorch containers to ensure environment consistency and reproducibility. The workflow includes automated preflight checks, live research of model-specific recipes, smoke testing on real data, and optional integration with HuggingFace Hub for model sharing. It is designed to be the fallback for models not covered by dedicated TAO skills, ensuring a rigorous, documented path from inspection to deployment.

What it does

  • Inspects and qualifies models and datasets for compatibility.
  • Audits local hardware and selects the appropriate NGC PyTorch container.
  • Performs live research to fetch model-specific training recipes and configurations.
  • Generates a complete, reproducible project structure including training and evaluation scripts.
  • Executes a mandatory smoke test on real data before full-scale training.
  • Pushes fine-tuned models to the HuggingFace Hub and emits a self-contained rerun skill.

When to use it

  • When you need to fine-tune a HuggingFace model and no dedicated TAO skill exists for it.
  • When you require a reproducible, containerized training pipeline on local NVIDIA hardware.
  • When you need to smoke-test a model locally before scaling to larger infrastructure.
  • When you want to automate the generation of model cards and training reports.

Advantages

  • Ensures environment consistency using NVIDIA-optimized NGC containers.
  • Prevents common training failures through mandatory smoke testing and preflight checks.
  • Reduces manual configuration by automating script generation and recipe research.
  • Provides a clear, auditable path from initial inspection to final model deployment.

Usage examples

“Fine-tune google/vit-base-patch16-224 on my local image dataset.”
The agent will inspect the model and dataset, verify hardware, generate the training project, run a smoke test, and initiate the fine-tuning process.
“Train a Llama-3-8B model using LoRA on my local GPU.”
The agent will research the appropriate LoRA configuration, set up the NGC environment, and execute the training pipeline with the specified parameters.

Install

  1. Download the skill as a ZIP:

    https://autogineer.com/agent-skills/nvidia/tao-finetune-huggingface-model/download

  2. Unzip it into the cross-client skills folder of your project:

    mkdir -p .agents/skills && unzip tao-finetune-huggingface-model.zip -d .agents/skills

Documentation by NVIDIA Corporation

<!-- Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. Licensed under the Apache License, Version 2.0; see http://www.apache.org/licenses/LICENSE-2.0 -->

tao-finetune-huggingface-model

Standalone install? If this session was not initialized by the TAO skill bank plugin, run the tao-setup skill first (host preflight, credentials, cross-skill discovery).

Local NVIDIA GPU fine-tuning for HuggingFace models, grounded in live-fetched documentation with curated references as a fallback safety net. One NGC container, a few focused scripts, one push to HF Hub. Follow the rules in this file; don't improvise.

Dedicated-model routing gate

Before Step 1 or any probe, image selection, package install, venv creation, or training-code generation, resolve model_id against the packaged model-owner registry. Use the absolute skill-bank root from which this file was loaded:

python <bank-root>/scripts/resolve_tao_model.py \
  --skill-bank <bank-root> \
  --model "$MODEL_ID" \
  --format json

The resolver matches model metadata, including huggingface_model_ids, network_arch, skill names, and legacy aliases. Routing is internal: a model ID and task are enough. Never require prompt boilerplate about skills, containers, or checkpoint formats.

  • Exit 0: stop this workflow and follow the owning model skill's environment, action metadata, preflight, and checkpoint preparation.

  • Exit 3: no packaged model skill owns the ID. This is the only result that permits Step 1 of the generic workflow.

  • Any other nonzero exit: ownership discovery is broken or ambiguous. Stop and resolve that error; do not silently fall back to generic Hugging Face training.

Hugging Face hosting never overrides ownership. Do not use this workflow to bypass a matched skill or ask the user to prescribe its internal preparation. For example, nvidia/Cosmos3-Nano routes to tao-finetune-cosmos-reason.

Do not create a host training venv in this workflow. Its default execution path is the NGC container documented below; any venv-based training path requires an explicit user request.

Order of authority (highest first):

  1. User input — explicit model_id, dataset_id, training_method, config.yaml overrides.

  2. Live research — model card, HF repo example, author finetune script, HF task docs, paper; always fetched (Step 3 + references/research-priorities.md).

  3. Curated references (references/*.md) — fallback when live research is silent/ambiguous.

  4. Your training-data memory — last resort; suspect, cross-check against (2)/(3).

Conflict resolution between (2) and (3) and the source-line discrepancy note are in references/research-priorities.md.

Inputs

Required:

  • model_id — HuggingFace model ID, e.g. google/vit-base-patch16-224

Conditional credentials (read from the session environment — exported before launching or sourced from a user-approved env file):

  • HF_TOKEN — only when the model/dataset is gated (read) or push_to_hub is on (write); public + public + push_to_hub: false needs none. Value never read — presence-only via [ -n "$HF_TOKEN" ].

  • WANDB_API_KEY, WANDB_PROJECT — only when WandB is enabled; WANDB_MODE=disabled opts out.

Dataset — exactly one:

  • dataset_id — HuggingFace dataset ID (source: hf)

  • local_dataset_path — local folder or file (source: local); optional local_dataset_format ∈ {auto, imagefolder, coco, voc, jsonl, arrow, parquet, csv} (default: auto-detect).

  • (omit) — agent recommends popular datasets (source: recommend)

Optional (have defaults):

  • task_type — auto-detected from config + model card

  • n_train=10000, n_eval=1000, n_epochs=3, lora_r=16

  • output_dir=./output/<model_short_name>

  • hf_model_repo — push target; if unset and HF_TOKEN has write access, auto-derived as <whoami>/<model_short_name>-finetuned.

  • push_to_hub=True — set to False to skip

  • skip_baseline=False — skip zero-shot baseline eval

Optional deliverables (off by default):

emit_progress_log: false   # output_dir/PROGRESS.md (per-step journal)
emit_report:       false   # reports/report.{pdf,html} with curves & samples
emit_unit_tests:   false   # tests/ with fake-data heterogeneous-batch tests

All values live in output_dir/config.yaml. Never hardcode in Python.

Execution platform

This skill orchestrates what to run; the platform skills own how to run it on a GPU host — read them first.

ConcernAuthoritative skill
GPU host runtime (driver 580, CUDA Toolkit 13.0, NVIDIA Container Toolkit 1.19.0)tao-skill-bank:tao-setup-nvidia-gpu-host
docker run flags, NGC auth, mounts, env passthrough, local/remote Docker job preflight (daemon, GPU smoke)tao-skill-bank:tao-run-on-docker

Default platform: local-docker — build a one-off image (run-<short>:latest) and run it on the local Docker daemon. Ask only when the user explicitly needs a different backend (Brev remote GPU, SLURM/Kubernetes); then run that platform's Preflight first and route the Steps 4–5 docker run commands through it. The GPU-runtime and presence-only credential preflights (values never read), the canonical docker run flag set, discovery of the execution platforms from the installed platform skills (tao-run-on-docker / -slurm / -kubernetes / -brev, plus any external one; on a runtime that surfaces only the core router skills, read skills/platform/tao-run-on-*/SKILL.md frontmatter), and the workflow-specific flags (--entrypoint /bin/bash -lc, PYTORCH_CUDA_ALLOC_CONF, --name hft_train) are in references/workflow-intake-preflight.md.

References — fallback safety net

Consulted only when live research is silent, ambiguous, or unavailable; live docs always win for the specific model and current API. Each step links the references it needs; full catalog in references/detailed-workflow.md.

Always-on: core-rules.md, error-playbook.md, compat-workarounds.md, model-discovery.md, dataset-recommendations.md, dataset-sources.md, dataset-patterns.md, hardware-container.md, research-priorities.md, cv-scripts.md, vlm-scripts.md, docker-runs.md, hub-push.md, pipeline-skill-template.md, deliverables.md. Opt-in (when their flag/need applies): progress-tracking.md, testing.md, reporting.md, workflow-intake-preflight.md, workflow-generate-train.md, workflow-push-rerun.md.

Rule: before falling back, log the live source you tried and why it was insufficient (config.yaml notes:, and PROGRESS.md if enabled). [FETCH LIVE] markers in cv-scripts.md / vlm-scripts.md are a research checklist, not code to inline — refetch the listed URL if a block has no Step 3 finding.

Core rules

Non-negotiable behaviors. Short version (full enumeration — hallucinated-imports list, never-without-approval list, full error-recovery and hardware-sizing tables — in references/core-rules.md, consult before any training-time decision):

  • Your HF-library knowledge is outdated. Fetch live docs (model card, HF repo example, task doc) before writing any ML code — don't generate trainer args / collator / transforms from memory (Step 3).

  • Smoke-test on real data with --max_steps 1 before any full run; no batch launches without a verified smoke.

  • Never silently substitute model_id, dataset_id, or training_method — if what the user asked for doesn't load, stop and ask.

  • Error recovery is minimal-change. OOM → halve batch, double grad_accum, enable gradient checkpointing (no LoRA switch without approval); NaN → reduce LR 10×; flat loss → inspect collator; same error 3× → stop and ask. Don't loop.

  • Dataset columns verified BEFORE the collator — rename in prepare_data.py; restructuring needed → stop and ask.

  • Hardware-sizing thumb (bf16): ≤3B → 24 GB, 7–13B → 80 GB, 30B+ → multi-GPU or LoRA on 1× 80 GB, 70B+ → 8× 80 GB or LoRA. Full finetune won't fit and no LoRA requested → ask before switching.

Workflow — 6 steps

Single pass, sequential; each step has a clear gate before the next begins.

Step 1 — Inspect & qualify

Goal: decide whether to proceed. Probe model + dataset, apply accept/reject, register applicable compat fixes, write the initial config.yaml.

Prerequisites: MODEL_ID, optional DATASET_ID / local_dataset_path, optional HF_TOKEN, OUTPUT_DIR (default ./output/<model_short_name>). Probes run in a CPU-only python:3.12-slim Docker container (bind-mounted .probe/ scratch) so the host needs no virtualenv — Docker must exist first. Docker-presence guard, container env, full probe invocation, and the model/dataset probe scripts are in references/workflow-intake-preflight.md, references/model-discovery.md, and references/dataset-sources.md.

Probe requirements:

  • Model: load AutoConfig, read model-card tags, detect task from architectures + tags + card examples (fallback logging in model-discovery.md).

  • Dataset: for recommended datasets, first present 3-5 choices from dataset-recommendations.md; for local data, bind-mount read-only and use dataset-sources.md format detection.

  • Reject early if the model config fails, the task is out of scope, no recipe source exists, or the dataset cannot load / match the task schema.

  • Evaluate compat-workarounds.md against the model/task; defer hardware-dependent rules to Step 2.

Write the initial config.yaml (model_id, task, dataset_id or local_dataset_path, research_sources: [] filled in Step 3, applicable_workarounds: from Step 1, notes: [] for reference fallbacks, push_to_hub: true default — annotated template in references/workflow-intake-preflight.md). Optionally rm -rf "$OUTPUT_DIR/.probe" once the gate is met.

Gate: config.yaml exists with model, dataset, task, applicable_workarounds; do not proceed if any field is missing.

Step 2 — Hardware audit & NGC image

Goal: verify Docker + GPU + disk, pick the NGC PyTorch image live, finalize hardware-dependent compat rules.

2a. Audit (hard gate) — three checks (commands in references/workflow-intake-preflight.md):

  1. GPU host runtime — tao-setup-nvidia-gpu-host's setup-nvidia-gpu-host.sh --backend docker --check-only; on fail, ask approval then re-run with --install --yes.

  2. Free-disk soft-warn — override via MIN_DISK_GB (default 100 GB); recommend ≥ 100 GB for NGC base (~20 GB) + HF cache + checkpoints + data.

  3. Conditional credential presence (values never read) — HF_TOKEN only when gated or push_to_hub is on; WANDB_* only when WandB is on.

Do not proceed to Step 4 on a hard-fail — Step 4's docker build pulls a 20+ GB NGC base, and a missing nvidia-container-toolkit only surfaces later as could not select device driver "" with capabilities: [[gpu]]. Record gpu_count, gpu_name, driver_major, vram_gb_per_gpu in config.yaml.

2b. Pick NGC image (live): from the NVIDIA Deep Learning Frameworks support matrix (https://docs.nvidia.com/deeplearning/frameworks/support-matrix/index.html), PyTorch NGC container section, pick the highest-versioned image where Min driver ≤ detected driver_major and container CUDA ≤ host CUDA Toolkit (match closely so cuDNN / TensorRT line up). Do not reject an image for an aN/bN/rcN PyTorch tag — NGC validates the full image; pick the newest CUDA-aligned one and let compat-workarounds.md handle per-version issues. If the matrix is unreachable, use the fallbacks in references/hardware-container.md; default nvcr.io/nvidia/pytorch:24.09-py3 <!-- unpinned: documented fallback --> (driver ≥ 545; SDPA+GQA bug — if num_key_value_heads < num_attention_heads, set attn_implementation: "eager"). Record ngc_image in config.yaml.

2c. Re-evaluate hardware-dependent compat rules: re-run the compat-workarounds.md walk for entries whose detect needs hw; update applicable_workarounds: in place.

2d. Model-fit check: estimate param_bytes ≈ 2×param_count (bf16); if

60% of vram_gb_per_gpu × 1e9, recommend LoRA in the user-facing summary.

Gate: config.yaml has ngc_image, gpu_count, gpu_name, driver_major, vram_gb_per_gpu; hardware-dependent compat fixes recorded.

Step 3 — Research the recipe

Goal: fetch the live recipe — training-data knowledge of transformers/trl/peft is suspect, so Step 3 is non-negotiable. Walk references/research-priorities.md in priority order (Priority 1 → 6); stop once you have, for the detected task:

  • AutoModel / processor class

  • Train + eval transforms

  • Collator

  • compute_metrics

  • Hyperparameter hints (LR, batch size, epochs, scheduler)

Record findings in meta/recipe.md, append source URLs to config.yaml: research_sources:. A slot with no live finding falls back to the matching scaffold (cv-scripts.md / vlm-scripts.md), logged as "fallback to scaffold — no live source for <slot>" under notes:. Conflict-resolution rules are in references/research-priorities.md.

Gate: every required slot filled, with a source URL or scaffold-fallback note.

Step 4 — Generate project & smoke-test

Goal: write all scripts, build the image, prepare data, run a 1-step smoke on real data (one docker build, two docker runs).

4a. Generate project files in output_dir/: config.yaml, Dockerfile, requirements.txt, prepare_data.py, train.py, run_eval.py, infer.py, optional merge_lora.py, optional tests/, .gitignore. Live Step 3 research is authority; cv-scripts.md / vlm-scripts.md give scaffold shape only. Apply every applicable_workarounds entry as a Dockerfile block, requirement pin, config override, or runtime env var. Hard rules: run_eval.py keeps that exact filename (avoids colliding with the HF evaluate package); every generated .py starts with the NVIDIA Apache-2.0 copyright header and any emitter fails when it is missing; emit_unit_tests: true generates and runs tests per references/testing.md. Script bodies, Dockerfile shape, and the emitter contract are in references/workflow-generate-train.md.

4b. Build, prepare, smoke — docker build -t run-<short>:latest ., then prepare_data and the --smoke --max_steps 1 run (references/docker-runs.md §1-3). Smoke pass criteria (in logs/smoke.log):

  • No exception

  • Loss is finite (not 0.0, not NaN)

  • grad_norm > 0 at step 1

If emit_unit_tests: true, also run pytest tests/ in the container. Any failure → STOP.

4c. Preflight summary — before full training, print and verify: reference URL, dataset columns, Hub target, monitoring target, NGC image, hardware, smoke loss/grad norm.

Gate: project files written, image built, smoke PASSED, preflight has no blank fields.

Step 5 — Train, evaluate, infer

Goal: baseline eval, full training, post-train eval, optional LoRA merge, 5 inference samples (all commands: references/docker-runs.md §4-8).

Sub-stepdocker-runs.mdSkip if
5a. Baseline eval (zero-shot)§4skip_baseline: true
5b. Full training (detached)§5—
5c. LoRA merge§6not VLM+LoRA
5d. Post-train eval§7—
5e. Inference (5 samples)§8—

Multi-GPU: prepend torchrun --nproc_per_node=$gpu_count to python train.py.

While training streams, watch docker logs -f hft_train: loss should drop within 10-20 steps; flat loss (collator/label-masking bug), NaN (LR too high), and OOM all stop the run — recovery in references/core-rules.md. If emit_report: true, run report.py after Step 5e per references/reporting.md.

Gate: all of:

  • checkpoints/final/ (or checkpoints/merged/ for LoRA) exists

  • reports/eval_results.json has a numeric primary metric

  • reports/baseline_results.json exists (unless skipped)

  • reports/inference_samples/ has 5 samples

  • wandb URL shows descending loss

Step 6 — Push & emit rerun skill

Goal: publish the run and make it reproducible without re-research.

Push per references/hub-push.md (weights, model card, eval/baseline JSONs, config.yaml, Dockerfile, requirements.txt, inference samples, reports when emitted) unless push_to_hub: false is explicit. Emit <output_dir>/skills/run-<short>/SKILL.md from references/pipeline-skill-template.md — substitute every placeholder, include full YAML metadata + the NVIDIA copyright HTML comment, and make any emitter fail if those are missing.

Gate (Done criteria): all of:

  • Step 5 gate met

  • HF Hub repo exists at the resolved URL with weights + card + results/ (unless push_to_hub: false)

  • <output_dir>/skills/run-<short>/SKILL.md exists, no <placeholder> left, with metadata + copyright HTML comment per pipeline-skill-template.md

Final message: wandb URL, HF Hub URL, baseline -> fine-tuned primary metric, reports/inference_samples/, and the rerun skill path.

Error playbook

On a known runtime error, consult the symptom → minimal-fix table in references/error-playbook.md (NGC entrypoint, PyTorch/Transformers regressions, numpy ABI, Albumentations bbox, PEFT/checkpointing, LoRA target breadth, CV augmentation gaps, OOM at step 0) before redesigning anything. When a row there fires twice across runs, lift it into compat-workarounds.md with a detect rule — auto-applied in Step 1 before the error can fire.

Communication style

  • Terse. No filler, no restating the request; one-word answers when appropriate.

  • Always include direct Hub and wandb URLs when referencing artifacts.

  • On error: state what went wrong, why, what you changed — no menus.

  • Never present "Option A/B/C" for a request with a clear answer. Act.

Files

references/testing.md

<!-- Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. -->

Unit Testing Reference (Phase 4.5)

Use this reference only when the parent SKILL.md points here for the current task. If this file conflicts with current SKILL.md, skill_info.yaml, schemas, or platform/model skills, the current authoritative source wins.

Contents

  • Philosophy

  • Generated test files

  • conftest.py Template (shared fixtures)

  • test_dataset.py Template (CV)

  • test_dataset.py Template (VLM)

  • test_collator.py Template (CRITICAL for VLMs)

  • test_model.py Template

  • test_smoke.py Template (1-step training smoke)

  • Running Tests — Phase 4.5 Command

  • What the tests would have caught in past runs

  • Task-branch-specific fixtures

Used in Phase 4.5 of tao-finetune-huggingface-model skill. Unit tests run against fake data inside the container before any GPU training. They catch shape/dtype/collation bugs that would otherwise only surface 5 minutes into a 2-hour training run.

Why this phase is mandatory: this is exactly where the VLM pipeline caught a regression in a prior run — pixel_values stacking failed because Idefics3 produces variable num_images per sample depending on image resolution. A 30-second collator test with two different-sized fake images would have caught it before the 15-minute wheel rebuild + training cycle.

Philosophy

  • Fake data, real code paths. Use PIL.Image.new() to create random RGB images, synthesize questions/answers, label arrays, and bounding boxes programmatically. Exercise the real Dataset, collator, model forward pass.

  • Variable, not uniform. Always include at least 2 samples with different shapes: different image sizes, different text lengths, different bbox counts. Catches the bugs uniform-batch testing misses.

  • Small and fast. Each test under 10 seconds. Full suite under 60 seconds.

  • Run inside the same container as training. Same PyTorch version, same transformers version.

Tests live under tests/. They are NOT shipped in the wheel — they are developer/CI artifacts.

Generated test files

FileWhat it tests
tests/conftest.pypytest fixtures: fake image(s), fake sample, fake batch, processor
tests/test_dataset.py__getitem__ returns correct keys, shapes, dtypes; label masking is non-trivial for VLMs
tests/test_collator.pyHeterogeneous batch collation works (critical — caught the VLM bug)
tests/test_model.pyModel loads; forward pass with fake batch returns finite loss
tests/test_smoke.py1 optimizer step on 2 fake samples updates weights

conftest.py Template (shared fixtures)

"""Shared pytest fixtures — fake data for all tests."""
import io
import random
import pytest
import torch
import yaml
from PIL import Image


@pytest.fixture(scope="session")
def cfg():
    return yaml.safe_load(open("config.yaml"))


def _fake_image(w: int = 224, h: int = 224, seed: int = 0) -> Image.Image:
    rng = random.Random(seed)
    arr = bytes(rng.randrange(256) for _ in range(w * h * 3))
    return Image.frombytes("RGB", (w, h), arr)


@pytest.fixture
def fake_image_small():
    return _fake_image(224, 224, seed=1)


@pytest.fixture
def fake_image_large():
    """Different size — important for VLM Idefics3 which splits high-res into tiles."""
    return _fake_image(512, 384, seed=2)


@pytest.fixture
def fake_cv_sample_classification(fake_image_small):
    """CV classification sample."""
    return {"image": fake_image_small, "labels": 0}


@pytest.fixture
def fake_cv_sample_detection(fake_image_small):
    """CV detection sample — 2 bboxes (small image)."""
    return {
        "image": fake_image_small,
        "objects": {
            "bbox": [[10.0, 20.0, 50.0, 60.0], [80.0, 80.0, 40.0, 40.0]],     # xywh
            "category_id": [0, 1],
            "area": [50*60, 40*40],
            "iscrowd": [0, 0],
        },
    }


@pytest.fixture
def fake_cv_sample_detection_large(fake_image_large):
    """CV detection sample — 4 bboxes on a DIFFERENT-SIZED image.

    Pairing this with `fake_cv_sample_detection` produces a heterogeneous
    batch (different image sizes AND different bbox counts) — the only kind
    that exposes detection-collator stacking bugs and label-list bugs.
    """
    return {
        "image": fake_image_large,
        "objects": {
            "bbox": [
                [5.0, 5.0, 30.0, 40.0],
                [50.0, 60.0, 80.0, 100.0],
                [200.0, 150.0, 60.0, 60.0],
                [300.0, 250.0, 50.0, 70.0],
            ],
            "category_id": [0, 1, 0, 2],
            "area": [30*40, 80*100, 60*60, 50*70],
            "iscrowd": [0, 0, 0, 0],
        },
    }


@pytest.fixture
def fake_cv_detection_batch(fake_cv_sample_detection, fake_cv_sample_detection_large):
    """CRITICAL: detection batch with DIFFERENT image sizes AND different bbox counts.
    Catches `torch.stack` collator bugs and `squeeze(0)` label-dict bugs."""
    return [fake_cv_sample_detection, fake_cv_sample_detection_large]


@pytest.fixture
def fake_vlm_sample_short(fake_image_small):
    """VLM sample with short question and short answer."""
    return {
        "image": fake_image_small,
        "question": "What color?",
        "multiple_choice_answer": "red",
    }


@pytest.fixture
def fake_vlm_sample_long(fake_image_large):
    """VLM sample with larger image and longer text — heterogeneity for collator."""
    return {
        "image": fake_image_large,
        "question": "Describe this scene in detail, including all visible objects and colors.",
        "multiple_choice_answer": "a colorful scene with many objects",
    }


@pytest.fixture
def fake_vlm_batch(fake_vlm_sample_short, fake_vlm_sample_long):
    """CRITICAL: batch with DIFFERENT image sizes — catches collator stacking bugs."""
    return [fake_vlm_sample_short, fake_vlm_sample_long]


@pytest.fixture(scope="session")
def processor(cfg):
    """Load the real processor for the configured model (once per session)."""
    import os
    from transformers import AutoProcessor, AutoImageProcessor
    token = os.environ.get("HF_TOKEN")
    try:
        return AutoProcessor.from_pretrained(cfg["model_id"], token=token)
    except Exception:
        return AutoImageProcessor.from_pretrained(cfg["model_id"], token=token)


@pytest.fixture(scope="session")
def tmp_arrow_dataset(tmp_path_factory, fake_vlm_sample_short, fake_vlm_sample_long):
    """Save a 2-sample Arrow dataset to disk (tests the load_from_disk path)."""
    from datasets import Dataset, Image as HFImage
    tmp = tmp_path_factory.mktemp("fake_data")
    ds = Dataset.from_list([fake_vlm_sample_short, fake_vlm_sample_long]).cast_column("image", HFImage())
    path = str(tmp / "mini")
    ds.save_to_disk(path)
    return path

test_dataset.py Template (CV)

"""Dataset.__getitem__ returns the right keys, shapes, and dtypes."""
import torch


def test_cv_classification_getitem(processor, fake_cv_sample_classification, tmp_path_factory):
    from datasets import Dataset, Image as HFImage
    from dataset import CVDataset

    tmp = tmp_path_factory.mktemp("cv_cls")
    ds_path = str(tmp / "mini")
    Dataset.from_list([fake_cv_sample_classification, fake_cv_sample_classification]) \
        .cast_column("image", HFImage()).save_to_disk(ds_path)

    ds = CVDataset(ds_path, processor)
    sample = ds[0]
    assert "pixel_values" in sample
    assert "labels" in sample
    assert sample["pixel_values"].ndim == 3
    assert sample["labels"].dtype == torch.long

test_dataset.py Template (VLM)

"""VLM dataset returns the right keys, labels are non-trivial."""
import torch


def test_vlm_getitem_shapes(processor, tmp_arrow_dataset, cfg):
    from dataset import VLMDataset
    ds = VLMDataset(tmp_arrow_dataset, processor, cfg)
    sample = ds[0]
    assert "input_ids" in sample
    assert "labels" in sample
    assert "pixel_values" in sample


def test_vlm_label_masking_non_trivial(processor, tmp_arrow_dataset, cfg):
    """Some (but not all) label tokens should be -100 — the prompt is masked, answer is not."""
    from dataset import VLMDataset
    ds = VLMDataset(tmp_arrow_dataset, processor, cfg)
    sample = ds[0]
    labels = sample["labels"]
    n_masked = (labels == -100).sum().item()
    n_unmasked = (labels != -100).sum().item()
    assert n_masked > 0, "Expected prompt tokens to be masked"
    assert n_unmasked > 0, "Expected answer tokens to NOT be masked (else loss will be 0)"

test_collator.py Template (CRITICAL for VLMs)

"""Collator must handle HETEROGENEOUS samples (different image sizes + text lengths).
This is where the pixel_values stacking bug bit us last time.
"""
import torch


def test_vlm_collator_heterogeneous_batch(processor, fake_vlm_batch, cfg, tmp_path_factory):
    """Samples with DIFFERENT image sizes must collate without 'list has no shape' error."""
    from datasets import Dataset, Image as HFImage
    from dataset import VLMDataset, collate_vlm

    tmp = tmp_path_factory.mktemp("hetero")
    ds_path = str(tmp / "mini")
    Dataset.from_list(fake_vlm_batch).cast_column("image", HFImage()).save_to_disk(ds_path)

    ds = VLMDataset(ds_path, processor, cfg)
    samples = [ds[0], ds[1]]

    # If this raises, production training will also crash
    batch = collate_vlm(samples, pad_token_id=processor.tokenizer.pad_token_id)

    # EVERY returned value must be a Tensor, not a list
    for k, v in batch.items():
        assert isinstance(v, torch.Tensor), f"{k!r} is {type(v).__name__}, expected Tensor"

    assert batch["input_ids"].shape[0] == 2, "Batch dim mismatch"
    assert batch["labels"].shape == batch["input_ids"].shape
    assert batch["pixel_values"].shape[0] == 2, "pixel_values batch dim mismatch"


def test_cv_detection_collator_heterogeneous_batch(processor, fake_cv_detection_batch, cfg, tmp_path_factory):
    """Detection: samples with DIFFERENT image sizes AND DIFFERENT bbox counts must
    collate without `stack expects equal size` and without 0-dim tensor errors in
    the labels dict."""
    from datasets import Dataset, Image as HFImage
    from dataset import CVDataset, make_collate_fn_detection

    tmp = tmp_path_factory.mktemp("hetero_det")
    ds_path = str(tmp / "mini")
    Dataset.from_list(fake_cv_detection_batch).cast_column("image", HFImage()).save_to_disk(ds_path)

    ds = CVDataset(ds_path, processor, task="object-detection", is_train=True)
    samples = [ds[0], ds[1]]

    # Each sample's labels must be dict-like (not a list-of-1-dict left over from
    # the processor). transformers 5.x returns BatchFeature (dict-like but not a
    # `dict` subclass) — assert by key membership instead of isinstance(dict).
    for i, s in enumerate(samples):
        assert "class_labels" in s["labels"], \
            f"sample {i}: missing class_labels in {type(s['labels']).__name__} — labels[0] extraction bug"
        # Class-label scalar tensors must keep shape (n_obj,), not be squeezed to 0-dim
        for k, v in s["labels"].items():
            if isinstance(v, torch.Tensor):
                assert v.ndim >= 1 or v.numel() <= 1, \
                    f"sample {i}: labels[{k!r}] has ndim={v.ndim} — squeeze(0) bug"

    collate_fn = make_collate_fn_detection(processor)
    batch = collate_fn(samples)

    assert "pixel_values" in batch
    assert isinstance(batch["pixel_values"], torch.Tensor), "pixel_values must be a Tensor (processor.pad)"
    assert batch["pixel_values"].shape[0] == 2, "batch dim mismatch"
    assert isinstance(batch["labels"], list) and len(batch["labels"]) == 2, \
        "labels must stay as list-of-dicts (variable bbox count per sample)"
    # Bbox counts differ between samples — verify the detection-specific shape
    assert batch["labels"][0]["class_labels"].shape[0] != batch["labels"][1]["class_labels"].shape[0], \
        "fixture should have different bbox counts; otherwise this test isn't exercising heterogeneity"

test_model.py Template

"""Model loads, forward pass produces finite loss on fake batch."""
import torch


def test_model_loads(cfg):
    from model import load_model_and_processor
    model, _ = load_model_and_processor(cfg)
    n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    assert n_trainable > 0, "Model has no trainable params"


def test_forward_pass_on_fake_batch(cfg, processor, fake_vlm_batch, tmp_path_factory):
    """End-to-end forward: dataset → collator → model.forward."""
    from datasets import Dataset, Image as HFImage
    from dataset import VLMDataset, collate_vlm
    from model import load_model_and_processor

    tmp = tmp_path_factory.mktemp("fwd")
    ds_path = str(tmp / "mini")
    Dataset.from_list(fake_vlm_batch).cast_column("image", HFImage()).save_to_disk(ds_path)

    model, _ = load_model_and_processor(cfg)
    ds = VLMDataset(ds_path, processor, cfg)
    batch = collate_vlm([ds[0], ds[1]], pad_token_id=processor.tokenizer.pad_token_id)
    batch = {k: v.to(model.device) for k, v in batch.items()}

    with torch.no_grad():
        out = model(**batch)

    assert torch.isfinite(out.loss), f"Loss is not finite: {out.loss}"

test_smoke.py Template (1-step training smoke)

"""One optimizer step on fake data. Catches misconfigurations that pass forward but fail backward."""
import torch


def test_one_training_step(cfg, processor, fake_vlm_batch, tmp_path_factory):
    from datasets import Dataset, Image as HFImage
    from dataset import VLMDataset, collate_vlm
    from model import load_model_and_processor

    tmp = tmp_path_factory.mktemp("smoke")
    ds_path = str(tmp / "mini")
    Dataset.from_list(fake_vlm_batch).cast_column("image", HFImage()).save_to_disk(ds_path)

    model, _ = load_model_and_processor(cfg)
    model.train()

    ds = VLMDataset(ds_path, processor, cfg)
    batch = collate_vlm([ds[0], ds[1]], pad_token_id=processor.tokenizer.pad_token_id)
    batch = {k: v.to(model.device) for k, v in batch.items()}

    optim = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=1e-4)
    out = model(**batch)
    out.loss.backward()

    # Every trainable param should have a gradient
    n_with_grad = sum(1 for p in model.parameters() if p.requires_grad and p.grad is not None)
    assert n_with_grad > 0, "No gradients computed — backward pass broken"

    optim.step()
    optim.zero_grad()


def test_trainer_one_step(cfg, processor, tmp_path_factory):
    """Run the actual HF Trainer for max_steps=1 on fake data.

    Why this matters: most "passes forward + 1 manual optim step" smoke tests
    miss bugs that live in `Trainer.training_step` itself — collator wiring,
    `remove_unused_columns` stripping label dicts, `bf16=True` cast paths,
    `log_history` shape (the summary entry has `train_loss` not `loss`).
    Running the real Trainer in unit tests catches those in seconds instead
    of in Phase 5.5 (which costs a full Docker rebuild to retry).

    Branch on task type to pick the right Dataset / collator / fixtures.
    Generate only the branch matching `cfg["task"]`.
    """
    from datasets import Dataset, Image as HFImage
    from transformers import TrainingArguments, Trainer
    from model import load_model_and_processor

    task = cfg["task"]
    tmp = tmp_path_factory.mktemp("trainer_smoke")
    ds_path = str(tmp / "mini")

    if task == "object-detection":
        from dataset import CVDataset, make_collate_fn_detection
        # Use the heterogeneous fixture — it's the most demanding shape
        # (different image sizes + different bbox counts) and the one most
        # likely to expose collator/label bugs in Trainer.training_step.
        from conftest import _fake_image  # noqa  (helper visible via conftest)
        # Build inline so this test doesn't need the batch fixture passed in
        samples = [
            {"image": _fake_image(224, 224, seed=1),
             "objects": {"bbox": [[10.0, 20.0, 50.0, 60.0]],
                         "category_id": [0], "area": [50*60], "iscrowd": [0]}},
            {"image": _fake_image(384, 256, seed=2),
             "objects": {"bbox": [[5.0, 5.0, 30.0, 40.0], [50.0, 60.0, 80.0, 100.0]],
                         "category_id": [0, 1], "area": [30*40, 80*100], "iscrowd": [0, 0]}},
        ]
        Dataset.from_list(samples).cast_column("image", HFImage()).save_to_disk(ds_path)
        train_ds = CVDataset(ds_path, processor, task="object-detection", is_train=True)
        model, _ = load_model_and_processor(cfg)
        collator = make_collate_fn_detection(processor)
    elif task == "image-classification":
        from dataset import CVDataset
        samples = [{"image": _fake_image(224, 224, seed=i), "labels": i % 2} for i in range(2)]
        Dataset.from_list(samples).cast_column("image", HFImage()).save_to_disk(ds_path)
        train_ds = CVDataset(ds_path, processor, task="image-classification", is_train=True)
        model, _ = load_model_and_processor(cfg)
        collator = None
    else:
        import pytest
        pytest.skip(f"trainer-step smoke not yet wired for task={task}")

    args = TrainingArguments(
        output_dir=str(tmp / "out"),
        max_steps=1,
        per_device_train_batch_size=2,
        learning_rate=1e-5,
        bf16=cfg.get("bf16", True),
        remove_unused_columns=False,
        report_to="none",
        logging_steps=1,
        save_strategy="no",
        eval_strategy="no",
        dataloader_pin_memory=False,
    )
    trainer = Trainer(model=model, args=args, train_dataset=train_ds, data_collator=collator)
    trainer.train()

    # Verify the same log-parsing pattern the production smoke uses
    step_log = next(
        (l for l in reversed(trainer.state.log_history) if "loss" in l), None
    )
    assert step_log is not None, "Trainer produced no step-level log entry"
    loss = step_log["loss"]
    assert loss == loss and loss != 0.0, f"smoke loss looks broken: {loss}"

Running Tests — Phase 4.5 Command

set -a; source /path/to/.env; set +a   # omit if already exported
docker run --rm --gpus all --shm-size=16g \
  -e HF_TOKEN \
  -e PYTHONUNBUFFERED=1 \
  -v $(pwd)/output_dir:/workspace \
  <ngc_image> \
  "cd /workspace && pip install -r requirements.txt pytest -q && \
   pip install dist/*.whl -q 2>/dev/null || python -m build --wheel --outdir dist/ -q && pip install dist/*.whl -q && \
   pytest tests/ -v --tb=short"

Add pytest>=7.0 to requirements.txt.

Gate: all tests pass. If any test fails, STOP and fix before Phase 5 wheel build.

What the tests would have caught in past runs

BugWhich test would catch it
pixel_values returned as list (heterogeneous images)test_vlm_collator_heterogeneous_batch
torch.stack fails on variable-sized detection imagestest_cv_detection_collator_heterogeneous_batch
squeeze(0) corrupts detection label dict (0-dim tensor)test_cv_detection_collator_heterogeneous_batch
All labels masked → loss always 0test_vlm_label_masking_non_trivial
ImportError from evaluate:main wheel entrywheel install step in Phase 4.5 runner
Wrong dtype= (fp32 fallback)test_model_loads trainable-param check
Forward pass NaNtest_forward_pass_on_fake_batch finite-loss check
Backward pass broken (detached graph)test_one_training_step grad check
remove_unused_columns strips label dicts inside Trainertest_trainer_one_step
log_history[-1] is summary, not step (train_loss vs loss)test_trainer_one_step
bf16=True × loaded-bf16 → optimizer underflow inside Trainertest_trainer_one_step (loss == 0 / NaN)

Without this phase, every one of these surfaced minutes or hours into real training.

Task-branch-specific fixtures

Generate only the fixtures and tests relevant to the detected task:

taskfixturestests
image-classificationfake_cv_sample_classificationtest_dataset.py (cls), test_model.py
object-detectionfake_cv_sample_detection (variable bbox count)test_collator.py (variable objects), test_model.py
semantic-segmentationfake_cv_sample_seg (image + mask)test_dataset.py, test_model.py
image-text-to-text (VLM)fake_vlm_sample_short + fake_vlm_sample_longall 4 test files — VLMs are the riskiest
text-generation (LLM)fake_text_sample_short + fake_text_sample_longtest_dataset.py, test_collator.py (variable length), test_model.py

Keep generated tests focused on the generated scripts — no speculative coverage.

Requirements and compatibility

Works with: Claude Code, Claude apps, OpenAI Codex, GitHub Copilot, Cursor, Gemini CLI, Kiro, OpenCode, Goose, Amp, Roo Code, JetBrains Junie, Factory Droid, Devin (ex-Windsurf), Mistral Vibe, OpenHands, Letta

“Works with” means the tool reads the open Agent Skills format; a skill written for one tool may use features others lack.

Needs: Declares allowed tools, Needs an API key, Needs Python

The author says: Requires docker + nvidia-container-toolkit, NVIDIA GPU (driver ≥ 545, ≥ 24 GB VRAM for ≤3B models), ~40 GB free disk. Optional credentials (read from the session environment) — HF_TOKEN is read only when the model/dataset is gated or `push_to_hub` is on; WANDB_API_KEY and WANDB_PROJECT only when WandB logging is enabled.

Security

What our scanner noticed· Medium risk

  • Large base64-encoded blobs
  • Instructions that use the network

A static scan is not a guarantee. Skills are third-party content: review a skill before you let an agent use it.

Attribution and licence

Created by
NVIDIA Corporation
Publisher
NVIDIA (Official)
Source
NVIDIA/skills · skills/tao-finetune-huggingface-model · commit fce3cca
Licence
Apache License 2.0 (redistribution allowed)
  • LICENSE (added by AutoGineer from LICENSE-APACHE)
  • ATTRIBUTION.md (written by AutoGineer)
Synced

AutoGineer is not affiliated with or endorsed by NVIDIA Corporation or NVIDIA.

Download ZIP

The ZIP holds the skill’s files, its licence files and an ATTRIBUTION.md. Download tao-finetune-huggingface-model.zip