| 名称 | tao-run-automl |
| 描述 | 运行基于容器的 NVIDIA TAO 网络 AutoML / 超参数优化(HPO),使用 AutoMLRunner。处理算法选择(bayesian、hyperband、asha、bohb、llm、hybrid、autoresearch)、WandB 实验跟踪、在任意 TAO SDK 平台上执行作业、结果解读以及按记录的自定义评估钩子。当用户提到 TAO AutoML、超参数优化、HPO、automl、automl_settings、AutoMLRunner、tao_automl、贝叶斯搜索、Hyperband、ASHA、LLM 引导搜索、autoresearch,或希望为任何 TAO 网络调整 train/distill/prune/quantize 动作参数时使用。模型动作默认使用模型技能解析出的容器镜像;基于 venv 的训练需要用户明确请求。平台无关——可在任何 SDK 上运行(Brev、SLURM、Kubernetes、Docker)。 |
| 开源协议 | Apache-2.0 compatibility: 需要 docker + nvidia-container-toolkit。工作流声明额外要求。 metadata: |
| 作者 | NVIDIA Corporation |
| 版本 | “0.1.1” allowed-tools: Read Bash Write tags: - automl - hpo - workflow - training - optimization - llm |
TAO AutoML
独立安装? 如果本会话尚未由 TAO 技能库插件初始化,请先运行
tao-setup技能(主机预检、凭据、跨技能发现)。
通过对 TAO 模型运行自动化超参数优化,该技能结合以下内容:
skills/models/<model_skill>/下的所选模型技能。skills/platform/<platform>/下的所选平台技能。AutoMLRunner,它生成建议、启动选定的动作作业、提取指标,并将结果反馈给优化器。
在模型元数据、平台预检、数据可见性、凭据、镜像选择和计算形状全部验证之前,不要启动。
执行运行时——硬性门槛
每个建议、基线评估、逐建议评估和最终评估都默认在所选模型动作解析后的 container_image 中运行。在任何训练环境设置之前,先从模型技能解析该镜像。本地检查点或 Hugging Face 模型 ID 不会改变此规则。
仅当用户明确要求 venv 执行时,才使用基于 venv 的模型动作执行。绝不要根据 local-docker、本地 GPU、已安装的 Python 环境或 pyproject.toml 的存在来推断 venv 模式。如果执行模式缺失,则为容器模式。为 tao_automl、TAO SDK 或平台适配器安装的主机/控制端 venv 仅用于控制平面;请明确说明这一点,并让所有子模型动作保留在解析后的容器镜像中。
参考映射
references/skill_info.yaml:此工作流的结构化元数据。- 明细参考文档:
automl-preflight-concepts.md用于前置条件和支持检查;automl-intent-algorithms.md用于搜索策略;automl-compression-literature.md用于蒸馏/剪枝/量化算法充分性以及未来压缩搜索路线图;automl-runner-configuration.md用于 runner/API/WandB 细节;automl-advanced-monitoring.md用于钩子、恢复和陷阱;而automl-examples.md用于对话示例。detailed-guide.md只是映射。 skills/models/<network>/SKILL.md:模型特定的数据集要求、指标、HPO 注意事项、检查点交接和已知失败。skills/models/<network>/references/skill_info.yaml:动作契约、容器镜像、输入、输出、上传排除项和mode。skills/platform/<platform>/SKILL.md:所选平台预检、凭据、资源形状、监控和取消。skills/core/tao-launch-workflow/SKILL.md:共享的平台、凭据、数据集可见性、镜像确认和用户确认输入模式。
预检
- 运行共享的启动输入检查。如果用户尚未选择平台,请询问;Brev、SLURM、Kubernetes 和 Docker 是平等选项。
- 在生成 runner 文件之前,运行所选平台技能的预检。
- 验证
nvidia-tao-automl可导入:
python -c "import tao_automl; from tao_automl.runner import AutoMLRunner; print('OK')"
如果缺失,显示来自 versions.yaml 的确切安装命令,并在安装前询问:
SB="${TAO_SKILL_BANK_PATH:-~/tao-skills-external}"
pip install "$($SB/scripts/resolve_versions_key.py wheels.tao_automl_<platform>)"
有效的平台 wheel 键为 tao_automl_brev、tao_automl_slurm、tao_automl_kubernetes、tao_automl_docker 和 tao_automl_all。仅在需要每个后端的开发机器上使用 all。仅当用户请求 LLM 引导算法时,才添加 ,llm。
模型支持门槛
每次运行前:
- 阅读模型
SKILL.md和references/skill_info.yaml。 - 确认模型的
automl_enabled: true,或模型技能明确将所选动作路由到 AutoML。 - 确认
<skill_dir>/schemas/<action>.schema.json存在且可解析。这是 AutoML 搜索空间的闸门。 - 对于非 TAO-Core 模型,如 Cosmos-RL 和 CLIP,还需要
references/spec_template_<action>.yaml;否则 runner 没有完整的动作默认值。 - 如果任何门槛失败,不要即兴设计搜索空间。报告缺失的包工件。
输入
在构造 runner 之前收集以下输入:
| 输入 | 要求 |
|---|---|
model_skill |
skills/models/ 下解析后的模型技能目录。在将用户别名如 network_arch 解析为打包的技能目录后,才接受它们。 |
network_arch |
从解析后的模型技能元数据中读取。 |
action |
要优化的动作,通常是 train、distill、prune 或 quantize。 |
platform |
一个受支持的 TAO 平台技能。 |
train_dataset / eval_dataset / 动作输入 |
使用模型特定的 spec 键和数据集布局。非训练动作通常还需要父检查点、教师检查点、校准数据或剪枝后的工件。 |
results_root |
适用于平台的本地、Lustre 或 S3 路径。 |
gpu_count、num_nodes |
遵守模型和平台限制。 |
container_image |
通过模型元数据和 versions.yaml 解析;向用户展示。 |
automl_algorithm |
默认 bayesian,除非用户要求其他算法或模型技能推荐其他算法。 |
metric、direction |
优先使用模型技能的验证/任务指标。 |
automl_budget |
按算法要求设置建议数量、最大 epoch/rung、并发数或种群大小。 |
永远不要询问秘密值。使用 [ -n "$VAR_NAME" ] && echo SET || echo UNSET 验证所需的环境变量。
启动前审查门槛
在启动任何建议作业之前,显示具体的启动审查并获得用户确认。此门槛适用于每个 AutoML 支持的模型/网络的每次 AutoML 运行;它并非 Cosmos 专属,也绝不能限定在单个模型技能范围内。即使平台和镜像预检已经通过,此条也适用。审查必须包括:
- 模型/网络、平台、镜像、GPU/节点形状和结果/工作空间根目录
- 数据集模式和具体的 spec 键,包括在可低成本读取时训练/评估样本数量
- 算法、预算、最大并发作业、指标和方向
- 可搜索参数及其范围,包括当用户未提供显式搜索空间时的默认值
- 为初始启动批次生成的确切推荐配置,在任何推荐作业提交前,在一个仅审查的步骤中生成
- 每个推荐的预计运行时间和总预期墙钟时间,并附上所用假设
- 预检后评估作业的自动基线和评估作业 ID、指标值和结果路径;如果模型没有可运行的 evaluate 动作或验证数据,则给出明确阻塞项
- AutoML 后对所选最佳检查点/模型的最终评估计划,包括指标、数据集和记录路径
如果估计时间长于用户声明的限制,或明显长于正常交互式运行,请在启动前询问是否减少建议数、epoch、数据集大小、验证频率或搜索空间。不要将多天估算隐藏在日志中。
自动基线评估作业
在平台、镜像、凭据、数据和模型预检通过后,在提交任何 AutoML 建议作业之前,先在选定的验证/评估数据上运行一次模型的 evaluate 动作。这是 AutoML 设置的必要部分,而不是对用户的可选“预训练评估”问题。使用 AutoML 训练运行开始时所用的同一基础模型或检查点、模型技能的 evaluate spec/模板,以及所选平台的正常作业提交路径。如果模型技能建议使用比训练更小的评估形状,请使用该形状并在启动审查中说明。
在请求确认启动建议之前,在启动审查中与用户分享评估指标数值。如果模型没有打包的 evaluate 动作、评估数据集缺失或评估作业失败,停止并报告阻塞项,而不是静默地回退到仅训练损失的 AutoML 运行。仅当用户明确接受运行将优化代理指标且没有影响基线时,才可在没有此基线的情况下继续。
AutoML runner 负责所选最佳检查点/模型的最终评估。当存在可运行的 evaluate 动作和验证/评估数据时,向 AutoMLRunner.run 传递 final_eval_fn(best_rec, train_job_id) 回调。该回调必须使用与基线相同的指标、数据集和方向评估所选最佳检查点/模型,在工作空间下存储结构化记录,并返回测量的指标值或包含 metric_value 及 record_path、job_id 等元数据的字典。不要在 runner.run 之后将最终评估作为代理侧步骤运行;返回的结果应包含 result["final_evaluation"] 及具体状态和原因。
依赖与数据预检
如果所选工作流需要对象存储或平台 CLI 且工具缺失,请报告缺失依赖并提供确切的安装命令,然后再继续。在用户批准后,重新运行 scripts/check_tao_launch_preflight.py 并带上 --install-missing-tools,以便安装最小所需包并立即重试路径验证。对于 S3 路径,在创建 runner 工件之前,先从启动平台验证凭据和路径可读性。不要等到第一个训练容器才发现缺少 AWS CLI、S3 客户端或 URI 不可读。
对于在每次训练试验期间读取大型媒体归档或目录的模型,请先在执行平台可见的存储上暂存或解压数据集一次,然后将所有推荐 spec 指向该暂存路径。在 <workspace>/evaluations/data_staging.json 中记录源 URI、暂存路径、可用时的字节/文件计数证据和时间戳。如果无法暂存,请在启动前审查中包含重复 S3 I/O 风险,并在将长 AutoML 预算用于该数据集之前询问。
当模型技能定义了样本数敏感约束时,在启动前强制执行这些约束。拒绝或限制任何会使所选数据集和 GPU 分片数产生零训练步的批量大小建议。在提交每个生成的推荐之前,使用 scripts/check_tao_launch_preflight.py --effective-batch-limit train_annotation=<batch_size>,<shard_count> 进行校验。如果某推荐随后因数据太小而无法容纳有效批量大小而失败,将其归类为无效配置,仅在剩余预算存在时替换或调整它,并在最终摘要中报告该修正。
当可从标注文件或廉价清单读取训练样本数量时,将其作为 automl_settings["train_sample_count"] 传给 AutoMLRunner.run,以便 runner 在提交作业前封顶不可能的建议,并将调整记录在 result["history"][i]["adjustments"] 中。
算法策略
| 算法 | 适用场景 | 必需调节旋钮 |
|---|---|---|
bayesian |
小/中预算、参数较少的默认选择。 | num_recommendations、metric、direction |
hyperband、asha |
配置多且早期 rung 便宜;ASHA 支持并行。 | max_epochs、reduction_factor、可选 max_concurrent |
bohb、dehb |
混合贝叶斯/进化搜索与多保真预算。 | 与 Hyperband 相同的 rung 预算字段 |
pbt |
长时间训练,调度应在训练中变异。 | 种群和世代预算 |
llm、hybrid、autoresearch |
用户明确要求 LLM 引导搜索并已配置端点。 | LLM 端点配置 + 预算 |
对于 distill,当蒸馏动作执行基于 epoch 的优化并写入检查点时,使用与训练类似的策略。对于单次 prune 和 quantize,默认使用 bayesian 或 bfbo,除非动作 schema/模型技能声明了类似 epoch 或校准预算的字段,使 hyperband/asha/bohb/dehb 有意义。当所选指标必须在压缩动作完成后由后续 evaluate/inference 动作计算时,使用 eval_fn。
优先采用模型技能的建议而非通用默认值。当模型技能表示启动、验证或检查点成本主导短试验时,避免使用 ASHA 或 Hyperband。
Spec 与搜索空间
将 spec 构建为嵌套字典。如果模型技能以点号表示法列出路径以便于阅读,则遍历路径并赋值到嵌套叶子节点;不要将扁平的带点字符串存储为 spec 键。
使用打包的所选动作 schema 来获取:
automl_default_parametersautoml_disabled_parameters- 有效的 min/max 范围
- 枚举、选项权重、条件、依赖项和热门参数
用户提供的搜索空间必须保持在 schema 约束内。对于具有离散选项的整数旋钮,如果模型技能明确要求,请包含 schema 所需的整数选项形状,而不是松散列表。
数据源覆盖是强制性的,除非模型技能表示启动器可以派生它们。当数据集使用直接标注/媒体路径时,保留用户提供的精确 spec 键。
指标策略
训练损失成本低但可能具有误导性。优先使用模型技能的任务指标。使用以下之一:
- 日志指标:
metric=<名称>、direction=maximize|minimize。 metric_extractor(logs, metric_name):当默认解析器有歧义时,解析模型日志。eval_fn(rec, train_job_id):当用户想要下游任务指标时,在每个推荐后运行模型的 evaluate 动作。
除非模型技能明确定义该映射,否则不要将 kpi 映射到指标。
对于每次 AutoML 运行,只要存在可运行的 evaluate 动作和验证/评估数据,就在预检之后、建议之前运行自动基线评估作业。最终报告必须比较基线指标、每个建议的指标和所选最佳指标,以便用户看到调优的影响。对于需要 eval_fn 来计算真实任务指标的模型技能,请使用该评估器而不是优化方便的损失,除非用户明确接受代理指标。
Runner 构造
仅在所选平台 SDK 的预检通过后使用它。构造 SDK 时不要将凭据嵌入代码中。
from pathlib import Path
from tao_automl.runner import AutoMLRunner
skill_bank = Path("<absolute-tao-skill-bank>")
model_skill = "<resolved-model-skill-directory>"
skill_dir = skill_bank / "skills" / "models" / model_skill
runner = AutoMLRunner(
skill_dir=str(skill_dir),
platform_sdk=sdk,
workspace_dir="<automl_workspace>",
)
result = runner.run(
automl_algorithm=algorithm,
automl_settings=automl_settings,
spec_overrides=spec_overrides,
automl_hyperparameters=automl_hyperparameters,
custom_param_ranges=custom_param_ranges,
metric_extractor=metric_extractor, # optional
eval_fn=eval_fn, # optional
final_eval_fn=final_eval_fn, # optional but required when final eval is runnable
)
仅当用户明确要求恢复、继续、恢复或检查现有实验时,才恢复现有工作空间。将普通的“运行 AutoML”请求视为全新运行。
监控
使用 runner 状态输出以及平台 SDK 的 get_job_status、get_job_logs 和 get_failure_analysis。对于活动作业,报告:
- 推荐 ID / 试验 ID
- 平台作业 ID
- 状态
- 当前指标
- 截至目前最佳指标
- 当前/最佳推荐所选超参数
- 经过时间,以及在有足够时序数据时的更新 ETA
失败时,判断是基础设施、数据可见性、镜像、凭据、spec/schema 还是模型代码失败。只修复最小原因,不要静默地在重复的无效推荐上花费额外预算。如果某个阻塞项在运行设置期间被修复,在显示更新后的预检/启动审查后继续原始任务,而不是让用户重新说明请求。
对于基于 LLM 的算法,在调用运行有效之前检查脑日志。验证 LLM 调用成功、生成了建议、先前指标被用于选择后续参数变更,并且日志显示保留/丢弃或等效的算法决策。如果大脑回退到随机采样,将 LLM 工作流分类为失败或阻塞,而不是将其视为有效的 LLM 引导运行。
结果交接
完成后:
- 按所选指标和方向识别最佳推荐。
- 返回最佳子作业 ID 及其结果路径。
- 使用模型技能的检查点/工件元数据和 SDK 帮助程序解析模型检查点或动作工件;不要猜测文件名(如
latest)。 - 报告准确的搜索空间、算法、预算、指标和平台。
- 报告自动基线评估作业 ID/结果路径/指标、所有推荐指标、最终评估状态/结果路径/指标、失败推荐及其根因、经过时间以及最终运行时说明。
- 如果该过程流入工作流(如 AutoML + DEFT),通过工作流声明的交接字段传递获胜 spec 覆盖和检查点。
- 按默认保留策略,验证支持清理、可安全剪枝的终止试验工件已被删除,并保留获胜训练工件。明确报告受保护提升/恢复父级或保守 Hybrid 结果。SDK 无法回收的远程绑定、命名卷或其他输出路由必须在第一个试验之前通过保留预检,而不是静默保留。
常见陷阱
- 不要期望在运行时存在
~/tao-core。schema 和模板必须打包在模型技能内。 - 不要从之前的运行推断数据集 URI。
- 不要预先计算 SDK 管理的输出路径;非 URI 输出值由 SDK 路由。
- 对于 SLURM,将大型数据集暂存在 Lustre 上,而不是在 GPU 分配时间内进行大量 S3 下载。
- 对于受门控的 HuggingFace 模型,验证
HF_TOKEN已设置,但不要读取它。 - 如果所有推荐都失败,停止并总结共享根因,而不是启动更多试验。
- 默认情况下不要禁用
automl_delete_intermediate_ckpt。保留每次试验可能消耗每个推荐一整套分布式检查点。 - 不要通过禁用清理来绕过保留预检失败,除非用户明确接受外部所有权以及对每个试验工件的手动生命周期管理。