nvidia-tao/tao-skills-bank/tao-run-automl

TAO AutoML

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

Run automated hyperparameter optimization for a TAO model by combining:

1. The selected model skill under skills/models/<model_skill>/.
2. The selected platform skill under skills/platform/<platform>/.
3. AutoMLRunner, which generates recommendations, launches selected action jobs,

extracts metrics, and feeds results back to the optimizer.

Do not launch until model metadata, platform preflight, data visibility, credentials, image choice, and compute shape are all proven.

Execution Runtime — Hard Gate

Every recommendation, baseline evaluation, per-recommendation evaluation, and final evaluation runs in the selected model action's resolved container_image by default. Resolve it from the model skill before any training-environment setup. A local checkpoint or Hugging Face model ID does not change this rule.

Use venv-based model execution only when explicitly requested. Never infer venv mode from local-docker, local GPUs, Python, or pyproject.toml. If absent, execution is container-backed. A host/controller venv for tao_automl, TAO SDK, or a platform adapter is control-plane-only; keep child model actions in the resolved container image.

Reference Map

• references/skill_info.yaml: this workflow's structured metadata.
• Split detailed references: automl-preflight-concepts.md for prerequisites

and support checks; automl-intent-algorithms.md for search policy; automl-compression-literature.md for distill/prune/quantize algorithm sufficiency and future compression-search roadmap; automl-runner-configuration.md for runner/API/WandB details; automl-advanced-monitoring.md for hooks, resume, and pitfalls; and automl-examples.md for conversation examples; and automl-common-pitfalls.md for recurring safety checks. detailed-guide.md is only the map.

• skills/models/<network>/SKILL.md: model-specific dataset requirements, metrics,

HPO notes, checkpoint handoff, and known failures.

• skills/models/<network>/references/skill_info.yaml: action contract,

container image, inputs, outputs, upload exclusions, and mode.

• skills/platform/<platform>/SKILL.md: selected platform preflight, credentials,

resource shape, monitoring, and cancellation.

• skills/core/tao-launch-workflow/SKILL.md: shared intake pattern for platform,

credentials, dataset visibility, image confirmation, and user confirmation.

Preflight

1. Run the shared launch intake. If the user has not chosen a platform, ask;

Brev, SLURM, Kubernetes, and Docker are equal peers.

1. Run the selected platform skill's preflight before generating runner files.
2. Verify nvidia-tao-automl imports:

python -c "import tao_automl; from tao_automl.runner import AutoMLRunner; print('OK')"

Then verify the selected platform's SDK constructs — importing tao_automl does not prove the platform backend is installed (e.g. DockerSDK() raises CredentialError without the docker package). See automl-preflight-concepts.md.

If missing, show the exact install command from versions.yaml and ask before installing:

SB="${TAO_SKILL_BANK_PATH:-~/tao-skill-bank}"
pip install "$($SB/scripts/resolve_versions_key.py wheels.tao_automl_<platform>)"

Valid platform wheel keys are tao_automl_brev, tao_automl_slurm, tao_automl_kubernetes, tao_automl_docker, and tao_automl_all. Use all only for development machines that need every backend. Add ,llm only when the user requests LLM-guided algorithms.

Model Support Gate

Before every run:

1. Read the model SKILL.md and references/skill_info.yaml.
2. Confirm automl_enabled: true for the model or that the model skill

explicitly routes the selected action to AutoML.

1. Confirm <skill_dir>/schemas/<action>.schema.json exists and parses. This

is the AutoML search-space gate.

1. For non-TAO-Core models such as Cosmos-RL and CLIP, also require

references/spec_template_<action>.yaml; otherwise the runner has no complete action defaults.

1. If any gate fails, do not improvise a search space. Report the missing

package artifact.

Inputs

Collect these before runner construction:

Never ask for secret values. Verify required env vars with [ -n "$VAR_NAME" ] && echo SET || echo UNSET.

Pre-Launch Review Gate

Before launching any recommendation jobs, show a concrete launch review and get user confirmation. This gate applies to every AutoML run for every AutoML-supported model/network; it is not Cosmos-specific and must not be scoped to a single model skill. This applies even when platform and image preflight already passed. The review must include:

• model/network, platform, image, GPU/node shape, and result/workspace root
• dataset mode and concrete spec keys, including train/eval sample counts when

they can be read cheaply

• algorithm, budget, max concurrent jobs, metric, and direction
• searchable parameters and ranges, including default values when the user did

not provide an explicit search space

• exact generated recommendation configs for the initial launch batch, produced

in a review-only step before any recommendation job is submitted

• estimated runtime per recommendation and total expected wall time, with the

assumptions used

• the automatic baseline eval job id, metric value, and result path from the

post-preflight eval job, or an explicit blocker if the model has no runnable evaluate action or validation data

• the post-AutoML final evaluation plan for the selected best checkpoint/model,

including metric, dataset, and record path

If the estimate is longer than the user's stated limit or materially longer than a normal interactive run, ask whether to reduce recommendations, epochs, dataset size, validation frequency, or search space before launch. Do not hide multi-day estimates in logs.

Automatic Baseline Eval Job

After platform, image, credential, data, and model preflight pass, run the model's evaluate action once on the selected validation/eval data before submitting any AutoML recommendation jobs. This is required AutoML setup, not an optional "pretrained eval" question for the user. Use the same base model or checkpoint that the AutoML training run starts from, the model skill's evaluate spec/template, and the selected platform's normal job submission path. If the model skill recommends a smaller shape for evaluation than training, use that shape and call it out in the launch review.

Share the eval metric number in the launch review before asking for confirmation. If a starting checkpoint exists but the baseline cannot be produced — no packaged evaluate action, missing eval dataset, failed eval job — stop and report the blocker instead of silently falling back to a training-loss-only run.

For training from scratch, record the baseline as unavailable and proceed; do not evaluate an empty checkpoint.

The runner owns final evaluation. When eval is runnable, pass final_eval_fn(best_rec, train_job_id) to AutoMLRunner.run; the result then carries result["final_evaluation"]. See automl-preflight-concepts.md for the callback, checkpoint, baseline, and from-scratch rules.

Dependency And Data Preflight

If the selected workflow needs object storage or a platform CLI and the tool is missing, report the missing dependency and offer the exact install command before continuing. After user approval, rerun scripts/check_tao_launch_preflight.py with --install-missing-tools so it installs the smallest needed package and immediately retries path verification. For S3 paths, verify both credentials and path readability from the launch platform before creating runner artifacts.

For models that read large media archives or directories during every training trial, stage or extract the dataset once to storage visible from the execution platform, then point all recommendation specs at that staged path. Record the source URI, staged path, byte/file-count evidence when available, and timestamp in <workspace>/evaluations/data_staging.json. If staging is not possible, include the repeated S3 I/O risk in the pre-launch review and ask before spending a long AutoML budget on it.

When the model skill defines sample-count-sensitive constraints, enforce them before launch. Reject or cap every batch-size recommendation that would create zero training steps for the selected dataset and GPU shard count. Use scripts/check_tao_launch_preflight.py --effective-batch-limit train_annotation=<batch_size>,<shard_count> for each generated recommendation before submitting it. If a recommendation later fails because the data is too small for the effective batch size, classify it as an invalid configuration, replace or adjust it only when remaining budget exists, and report the correction in the final summary. When train sample count is known from an annotation file or cheap manifest read, pass it as automl_settings["train_sample_count"] to AutoMLRunner.run so the runner can cap impossible recommendations before submitting a job and record the adjustment in result["history"][i]["adjustments"].

Algorithm Policy

For evaluate or inference, default to Bayesian/BFBO-style search over the selected action's prompt, decoding, preprocessing, or runtime config knobs. Use a task metric from the action outputs/logs and set direction explicitly when the metric name is ambiguous. Do not use training-loss assumptions for actions that do not update weights.