TAO-HF Integration Skill

Integrate a HuggingFace (HF) Computer Vision model into the NVIDIA TAO Toolkit ecosystem. Work the phases iteratively — not purely linearly — via a build → test → debug → fix → retest loop at every step: when something fails, diagnose and fix before moving on; when it passes, move to the next step.

This SKILL.md is the workflow coordinator. Each phase has a dedicated references/phase-N-*.md with the full step-by-step content, code, docker invocations, and gates. Read the matching reference at the start of each phase — the summaries below are not sufficient.

Local-Only Rule

All work is strictly local. Do NOT push/commit/branch on any remote (GitLab, GitHub, HuggingFace), create merge/pull requests or issues, or upload/publish Docker images to any registry or artifact store. You may only read/clone from remotes — all edits, Docker builds, and test runs stay on the local machine.

Submodule Override Strategy

The user clones the four TAO repos (tao-core, tao-pytorch, tao-deploy, tao-dataservices) independently into one working directory. The tao-core/ submodule nested inside each repo points to the original unmodified commit; modifications only exist in the top-level tao-core/. Always install from the top-level tao-core/, never <repo>/tao-core/ — the nested submodule silently ignores all modifications. Override rules: (1) mount the working directory -v $(pwd):/workspace; (2) pip install /workspace/tao-core FIRST, before tao-pytorch/tao-deploy; (3) PYTHONPATH top-level tao-core first, e.g. -e PYTHONPATH=/workspace/tao-core:/workspace/tao-pytorch. See references/cross-cutting.md for the directory tree.

Execution platform

Every test, smoke run, and end-to-end validation executes inside a locally prepared TAO Toolkit container (tao-pytorch-base:latest, tao-deploy-base:latest, optionally tao-dataservices-base:latest — all from Phase 0). The platform skills own how to run them; this skill specifies what. Default platform: local-docker. Phase 0 delegates the driver / CUDA / NCT preflight to tao-setup-nvidia-gpu-host. See references/cross-cutting.md for the authoritative-skill table, bind-mount rationale, and canonical docker-run flag set.

Phase Map

Cross-cutting refs: workflow-consistency.md (CLI flow, config field paths, cross-phase dependencies); cross-cutting.md (platform, isolation, module pitfalls, debugging).

IMPORTANT — Continuous Execution Through Phase 6: do NOT stop after Phases 3–5 to wait for the user to run tests. Phase 6 is mandatory — not complete until tests pass inside the containers and the end-to-end pipeline is validated.

Development Loop

At every step: write code → test immediately (import check, unit test, or dry-run) → if it fails, read traceback → diagnose → fix → retest; if it passes, move on. Do NOT accumulate untested code — testing only at the end compounds bugs.

Debugging Playbook

When something fails, consult the symptom → likely-cause → fix table in references/cross-cutting.md before trying random fixes — it covers ModuleNotFoundError, BACKBONE_REGISTRY KeyError, shape mismatch, NaN loss, ONNX/TRT build failures, TRT-vs-PyTorch accuracy gaps, OOM, DDP hangs, checkpoint load failures, and stale-submodule config issues.

Environment Isolation Strategy

All Python work runs inside Docker containers — no host venvs, no pip installs into host Python (the host needs only Docker, from tao-setup-nvidia-gpu-host). Three contexts: A (Phase 1 HF inspection in tao-hf-inspect, python:3.12-slim fallback), B (Phase 3/4/6 smoke/L0/e2e in the prepared container, source via pip install /workspace/tao-core && python setup.py develop), C (host-bind-mount scratch). See references/cross-cutting.md for the contexts in full and the four numbered rules verbatim (--check-only host packages; Phase 1 --user $(id -u):$(id -g) vs. root; HOME=/workspace/PIP_USER=1 fallback; distro package-manager list; root:root trade-off; tao-hf-inspect cleanup).

Phase 0 — Prerequisites Check

Goal: verify Python 3.10+ and git; delegate the driver / CUDA / Docker / NVIDIA Container Toolkit host check to tao-setup-nvidia-gpu-host; verify NGC docker login for nvcr.io. Then ask the user for the TAO Toolkit image references (tao-pytorch, tao-deploy, optionally tao-dataservices), pull, and prepare local tags tao-pytorch-base:latest, tao-deploy-base:latest, tao-dataservices-base:latest for later phases — preparation removes the pre-installed released TAO packages so the user’s /workspace/... clones install/load via pip install /workspace/tao-core && python setup.py develop. Hard stop on any failed check. Required user inputs: the image references + credentials (NGC login, HF_TOKEN). Full commands, prompt wording, and per-image Dockerfile snippets: the Phase 0 reference.

Gate: all prerequisite checks pass; the user supplied the required image references; tao-pytorch-base:latest and tao-deploy-base:latest exist locally; tao-dataservices-base:latest exists if dataservices work is anticipated.

Phase 1 — Information Gathering & Validation

Goal: decide whether to proceed at all. Gather credentials, locate/clone the four TAO repos, create a consistent working branch, launch the tao-hf-inspect container (Context A), validate the HF model is CV with a supported pipeline_tag, extract config + state-dict schema, sanity-check ONNX export, clean up. Full steps: Phase 1 references.

Reject if: pipeline_tag is NLP / audio / LLM (non-CV); AutoConfig raises; or ONNX export fundamentally cannot work (no rewrite path).

Gate: all 4 TAO repos located/cloned with a consistent branch; pipeline_tag confirmed CV; model_type, image_size, hidden_size, num_labels extracted; state-dict keys documented + HF→TAO remapping plan drafted; ONNX export sanity check passed (or failure understood); user confirmed model_short_name + task type. (Full checklist: phase-1-inspection.md.) Present findings and get user confirmation first.

Phase 2 — Codebase Exploration

Goal: find the closest existing TAO reference model for the detected pipeline_tag, read its implementation across tao-core / tao-pytorch / tao-deploy, and decide whether the backbone exists in backbone_v2/ or is new.

The HF pipeline_tag → TAO reference model mapping (classification → classification_pyt, detection → dino/rtdetr, segmentation → segformer, instance → mask2former, panoptic → oneformer, zero-shot → grounding_dino, depth → mono_depth) drives everything downstream (config, architecture, loss, ONNX shape, TRT builder, deploy classes, metrics, dataset format). See the Phase 2 references for the full reference list (12 files per model), the backbone_v2/ and tao-dataservices coverage checks, and per-task architecture.

If a new backbone is needed, decide the strategy (timm wrap > re-implement > HF black-box wrap) before Phase 3 — it changes weight loading, ONNX export, deploy. Never dual-inherit from transformers.PreTrainedModel and BackboneBase (metaclass conflict — compose instead).

Gate: reference TAO model identified + all 12 reference locations read; task-type implications understood (architecture, loss, ONNX outputs, deploy classes, metrics, dataset); backbone coverage decided (reuse / wrap timm / new); dataservices coverage checked. (Full checklist: phase-2-codebase.md.)

Phase 3 — TAO Core Configuration & Native Implementation

Goal: write the tao-core config schema + the tao-pytorch trainer / native inference / evaluation, smoke-testing between steps. (<model_name> = snake_case short-name; <ModelName> = PascalCase.)

Steps 1–7 (each builds on the previous, smoke-test between): tao-core config (1), tao-pytorch trainer (2), multi-GPU/multi-node (3), native inference → result.csv (4), native evaluation → results.json (5), MLOps for training and eval/infer → status.json (6–7). The ExperimentConfig(CommonExperimentConfig) must contain model, dataset, train, evaluate, inference, export, gen_trt_engine, quantize. All ??? fields are MISSING (user supplies via YAML/CLI); the augmentation.mean/std, model.head.in_channels, checkpoint-name, and onnx_file matches are in the checklist below.

Full per-step bodies, code, the canonical experiment_spec.yaml, and smoke-test commands: the Phase 3 references.

Gates: Step 1 — ExperimentConfig imports cleanly in-container; Step 2 — build_model(cfg) runs + PLModel instantiates in-container; Phase 3 — all 7 steps complete, smoke tests pass, no missing __init__.py.

Phase 4 — Export, Deployment & TensorRT Integration

Goal: ONNX export from tao-pytorch, then TRT engine builder + inference + evaluation in tao-deploy reusing the tao-core ExperimentConfig.

Steps 8–11: ONNX exporter (8 — task-specific input/output names, batch_size=-1 ⇒ dynamic batch); TRT engine builder (9 — subclass EngineBuilder or reuse ClassificationEngineBuilder; write specs/{gen_trt_engine,inference,evaluate}.yaml, same ExperimentConfig schema, augmentation.mean/std MUST match training); TRT inference → result.csv (10); TRT eval → results.json (11). See the Phase 4 reference for full code and the Phase 3+4 gate (3 in-container checks: imports, model build + forward, ONNX round-trip).

Module pitfalls: tao-pytorch and tao-deploy have separate hydra_runner and monitor_status — use the deploy versions in deploy scripts. ExperimentConfig comes from nvidia_tao_core in both (same schema/field paths).

Phase 3+4 gate: all three in-container checks pass (tao-pytorch imports + model + ONNX export; tao-deploy imports).

Phase 5 — Packaging & L0 Testing

Goal: register the model as a console_script in both repos and add unit tests.

Steps 12–15: register '<model_name>=...:main' in console_scripts of tao-pytorch/setup.py (12) and tao-deploy/setup.py (13, creating the deploy entrypoint/<model_name>.py via entrypoint_hydra); deploy L0 tests (14); trainer L0 tests — Trainer(..., fast_dev_run=True) + @pytest.mark.cv_unit @pytest.mark.<model_name> (15). See the Phase 5 reference for exact entry-point strings, code, and L0 test file lists.

Gate: entrypoints registered; pytest files exist and follow the marker convention. Do NOT stop — go directly to Phase 6.

Cross-Phase Data Flow & Consistency Verification

Before Docker testing, verify the chain train → export → gen_trt_engine → inference / evaluate (the *_model_latest.pth → .onnx → .engine artifact flow + the config fields each stage reads/writes — full diagram in references/cross-cutting.md).

Consistency checklist (verify before proceeding): self.checkpoint_filename → the *_latest.pth name evaluate.checkpoint / export.checkpoint reference; augmentation.mean/std identical across training spec, inference.yaml, evaluate.yaml, engine-builder preprocess_mode; ONNX input_names=['input'] / output_names=['output'] (detection/instance-seg use task-specific names); export.input_width/input_height match dataset.img_size; model.head.in_channels matches model_params_mapping.py; classes.txt at dataset.root_dir readable by both repos; all __init__.py exist (incl. scripts/__init__.py for get_subtasks() via pkgutil). Full paths: workflow-consistency.md.

Phase 6 — Container Testing & End-to-End Validation

Mandatory — start immediately after Phase 5. All TAO models ship as Docker images; code that only works outside a container is incomplete. Testing runs directly inside the TAO Toolkit container — no image build in the loop: mount → install source (setup.py develop) → run pytest / pylint / pydocstyle / flake8 directly. Use vanilla commands, NOT the ci/run_functional_tests.py / ci/run_static_tests.py wrappers (internal-mirror-only; public github.com/NVIDIA-TAO/ mirrors have no ci/ dir).

Steps 16–25: verify local image tags exist (16); unit tests for tao-core / tao-pytorch (-m cv_unit, --shm-size=16G) / tao-deploy (17–19); lint (20); wheels (21); end-to-end — train dry-run + export in one tao-pytorch session, then gen_trt_engine + inference + evaluate in one tao-deploy session (same session critical — --rm discards installs) (22); cross-check native vs TRT (23); debug shells (24); optional release images (25).

Full commands (every docker run, per-container env-vars, exact pytest / lint invocations + full-suite variants, the train/export/gen_trt_engine/inference/evaluate one-liner with all CLI overrides, the ci/ note, the fix-and-retest loop) and build scripts / runner patterns: see the Phase 6 references.

Phase 6 gate (Done criteria): tao-core / tao-pytorch / tao-deploy unit tests pass in their containers; static tests pass (or only legacy lint warnings); wheels build; end-to-end <model_name>_model_latest.pth → model.onnx → model.engine → non-empty result.csv + results.json; native vs TRT agree within tolerance.

Phase 7 — Optimization & Tuning (conditional)

Enter only if Phase 6 passes but accuracy / latency / size needs improvement. Ask the user for target metrics first.

Diagnostic categories: accuracy too low; TRT-vs-native gap; training too slow; inference too slow. Techniques: Step 27 — hyperparameter tuning; Step 28 — INT8 quantization (PTQ via torchao / modelopt, TRT INT8 + calibration); Step 29 — channel pruning + retrain; Step 30 — knowledge distillation; Step 31 — resolution tuning (TAO interpolates ViT positional embeddings automatically). See the Phase 7 reference for each category’s checks, config blocks, YAML overrides, decision tree, and rationale.

Argument

$ARGUMENTS

If provided, interpret $ARGUMENTS as the HuggingFace model ID or URL to start Phase 1. If credentials or model short-name are not included, ask the user for them before proceeding.