tao-finetune-cosmos-embed

Cosmos-Embed

Cosmos-Embed1 是一个视频与文本联合嵌入器，适用于文本到视频检索、视频到视频搜索、零样本/kNN 分类和语义去重。打包的 CLI 为 cosmos-embed1，支持 train、evaluate、inference 和 export。

容器镜像和各操作的命令位于 references/skill_info.yaml。精简的起始配置位于 references/spec_template_*.yaml。

训练操作策略

该模型在模型层面启用了 AutoML。在处理任何训练阶段请求之前，请先阅读 references/skill_info.yaml，并根据显式的 automl_policy 值或用户的工作流请求确定本次运行的覆盖设置。将“关闭 AutoML”“禁用 AutoML”“不进行 HPO”或“普通训练”等表述视为仅对此次运行设置 automl_policy: off；否则默认使用 auto。当 automl_policy: auto、automl_enabled: true，且已打包 schemas/train.schema.json 和 references/spec_template_train.yaml 时，默认使用此模型的 skill_dir，通过 tao-skill-bank:tao-run-automl 执行训练操作。保留工作流/应用对数据集、配置、输出目录、GPU/平台设置、父检查点和 automl_policy 的覆盖设置。只有在 automl_policy: off 或缺少已打包的训练 schema/模板时，才直接进行模型训练；如果缺少 schema，请说明 AutoML 已启用，但在生成 schema 之前无法用于此模型。

evaluate、inference、export 和部署流程等非训练操作仍在此模型技能中执行。每次运行的 automl_policy 覆盖设置不会更改模型元数据。

快速开始

使用 references/skill_info.yaml 中声明并通过 versions.yaml 确定的已发布 Cosmos-Embed 容器。正常使用技能时，请勿从私有的 Cosmos-Embed1 源代码树构建；只有在开发容器本身时才从源代码构建。

TAO_SKILL_BANK_PATH="${TAO_SKILL_BANK_PATH:-$PWD}"
COSMOS_EMBED_IMAGE="${COSMOS_EMBED_IMAGE:-$(
  python "$TAO_SKILL_BANK_PATH/scripts/resolve_tao_image.py" \
    --skill-bank "$TAO_SKILL_BANK_PATH" \
    --model cosmos-embed \
    --action train \
    --format json |
  python -c 'import json,sys; print(json.load(sys.stdin)["image"])'
)}"
docker pull "$COSMOS_EMBED_IMAGE"

预期的本地工作区布局：

workspace/
├── data/
│   ├── msrvtt_test_1k.json
│   └── video/
│       ├── video7020.mp4
│       └── ...
├── model/
│   └── Cosmos-Embed1-224p/        # optional if using HF repo id
├── specs/
│   ├── train.yaml
│   ├── evaluate.yaml
│   ├── inference.yaml
│   ├── export_onnx.yaml
│   └── export_hf.yaml
└── results/

除非本地 Docker/平台技能指定了更严格的环境专属命令，否则所有操作均使用以下 Docker 选项：

TAO_SKILL_BANK_PATH="${TAO_SKILL_BANK_PATH:-$PWD}"
COSMOS_EMBED_IMAGE="${COSMOS_EMBED_IMAGE:-$(
  python "$TAO_SKILL_BANK_PATH/scripts/resolve_tao_image.py" \
    --skill-bank "$TAO_SKILL_BANK_PATH" \
    --model cosmos-embed \
    --action train \
    --format json |
  python -c 'import json,sys; print(json.load(sys.stdin)["image"])'
)}"
RUN_ROOT="${RUN_ROOT:-$PWD}"
DOCKER_COMMON=(
  --rm --gpus all --ipc=host --network=host
  --shm-size=64g
  --ulimit memlock=-1
  --ulimit stack=67108864
  -e HF_TOKEN
  -v "$RUN_ROOT/data:/data:ro"
  -v "$RUN_ROOT/model:/model"
  -v "$RUN_ROOT/specs:/specs:ro"
  -v "$RUN_ROOT/results:/results"
)

训练：

docker run "${DOCKER_COMMON[@]}" "$COSMOS_EMBED_IMAGE" \
  cosmos-embed1 train -e /specs/train.yaml results_dir=/results

评估：

docker run "${DOCKER_COMMON[@]}" "$COSMOS_EMBED_IMAGE" \
  cosmos-embed1 evaluate -e /specs/evaluate.yaml results_dir=/results

推理：

docker run "${DOCKER_COMMON[@]}" "$COSMOS_EMBED_IMAGE" \
  cosmos-embed1 inference -e /specs/inference.yaml \
  'inference.query.input_texts=["a man is singing on stage"]' \
  inference.k=5 \
  results_dir=/results

导出 ONNX：

docker run "${DOCKER_COMMON[@]}" "$COSMOS_EMBED_IMAGE" \
  cosmos-embed1 export -e /specs/export_onnx.yaml \
  export.checkpoint=/results/train/cosmos_embed1_model_latest.pth \
  export.onnx_file=/results/export/cosmos_embed1_combined.onnx \
  results_dir=/results

导出 HuggingFace 格式：

docker run "${DOCKER_COMMON[@]}" "$COSMOS_EMBED_IMAGE" \
  cosmos-embed1 export -e /specs/export_hf.yaml \
  export.checkpoint=/results/train/cosmos_embed1_model_latest.pth \
  export.hf_output_dir=/results/export_hf/cosmos_embed1_hf \
  results_dir=/results

冒烟测试覆盖设置

进行小规模功能检查时，保留相同的配置，并覆盖开销较大的参数：

train.max_iter=1
train.validation_iter=2
train.checkpoint_iter=1
train.optim.optim=adamw
dataset.train_dataset.batch_size=1
dataset.val_dataset.batch_size=1
dataset.train_dataset.workers=0
dataset.val_dataset.workers=0

如果没有本地 Cosmos-Embed1 预训练检查点或 HuggingFace 令牌，请将 model.pretrained_model_path=null，以进行仅检查流程的冒烟训练。在此模式下，模型质量没有意义，但仍可测试 train/evaluate/inference/export 操作路径。

针对少量数据进行评估和推理冒烟测试：

evaluate.callbacks.embedding_visualization=false
evaluate.callbacks.max_eval_samples=8
dataset.test_dataset.batch_size=1
dataset.test_dataset.workers=0
inference.k=2
dataset.inference_dataset.batch_size=1
dataset.inference_dataset.workers=0

数据格式

MSR-VTT 路径需要本地视频 glob 和 JSON 元数据文件：

dataset:
  train_dataset:
    dataset_type: msrvtt
    mp4_urls: /data/video/*.mp4
    metadata: /data/msrvtt_test_1k.json

列表格式的元数据行至少必须包含 video 和 caption：

{"video_id": "video7020", "video": "video7020.mp4", "caption": "a woman creating a fondant baby and flower"}

数据集加载器会从本地 .mp4 文件名中提取视频 ID，并筛选出元数据中存在的视频。如果运行时未找到任何视频，请检查 mp4_urls 是否指向容器内的 glob，以及元数据中的 video 名称是否与文件名匹配。

模型权重

• 本地 HF 目录：将其挂载到 /model 下，并设置 model.pretrained_model_path=/model/Cosmos-Embed1-224p。
• HuggingFace 仓库：设置 model.pretrained_model_path=nvidia/Cosmos-Embed1-224p；如果访问受限，请传入 HF_TOKEN。
• 微调后的检查点：下游操作默认使用 /results/train/cosmos_embed1_model_latest.pth。

变体：

变体	分辨率	帧数	嵌入维度
Cosmos-Embed1-224p	224 x 224	8	256
Cosmos-Embed1-336p	336 x 336	8	768
Cosmos-Embed1-448p	448 x 448	8	768

确保 model.network.embed_dim、model.input_hw 和 model.network.spatial_resolution 与所选变体保持一致。

重要参数

参数	说明
train.num_gpus	单 GPU 时设为 1，设为 >1 会自动启动 torchrun，设为 -1 会自动检测可见 GPU。
train.max_iter	主要训练时长。仅在冒烟测试时使用 1。
train.optim.optim	支持时，fused_adamw 速度更快；冒烟测试和可移植性场景下，adamw 更稳妥。
model.lora.enabled	启用 LoRA。启用 LoRA 时，将 model.network.visual_encoder.transformer_engine=false。
model.lora.lora_rank	LoRA 秩。建议从 8 开始；手动调参或进行 AutoML 式搜索时，可尝试 4、8 或 16。
model.lora.lora_alpha	LoRA 缩放因子。建议从 16 开始；除非实验结果另有说明，否则保持在 lora_rank 的 2 倍左右。
model.lora.lora_dropout	LoRA dropout。建议从 0.1 开始；针对小型数据集，可尝试 0.0、0.05 和 0.1。
model.lora.bias	偏置策略：none、all 或 lora_only。除非有意训练偏置，否则保持为 none。
model.lora.use_rslora / use_dora	可选的 LoRA 变体。每次只启用一种，并将设置与检查点一并记录。
model.lora.target_modules	LoRA 注入时可选的模块名称模式。留空则使用默认的 ViT + Q-Former 注意力/MLP 目标模块。
model.lora.modules_to_save	可选模块，与 LoRA 一同保持完全可训练。除非需要保留任务专用的头部，否则留空。
evaluate.load_dataset_pkl / save_dataset_pkl	缓存评估嵌入。
inference.load_dataset_pkl / save_dataset_pkl	缓存搜索数据库，以便重复检索。
export.mode	video、text、combined 或 huggingface。
export.on_cpu	建议在 CPU 上导出，以避免设备不匹配问题。

LoRA 和 AutoML 说明

要进行参数高效微调，请设置 model.lora.enabled=true，并保持 model.network.visual_encoder.transformer_engine=false；TAO Core 的 Cosmos-Embed1 配置说明指出，PEFT 无法向 Transformer Engine 层注入适配器。请将上述 LoRA 字段视为手动调参或 AutoML 式搜索的首批候选参数， 然后再解冻更大的模型模块。除非用户明确需要自定义适配器放置位置， 否则不要更改 target_modules 或 modules_to_save。

S3 暂存

Cosmos-Embed1 CLI 使用本地路径和 Python glob，而不是原始的 s3://.../*.mp4 URI。对于由 S3 提供数据的运行，请先将部分或完整数据集暂存到执行主机/容器的文件系统中，然后在 spec 中使用 /data/video/*.mp4 等本地路径。

暂存 MSR-VTT 数据时，推荐使用以下 S3 布局：

s3://bucket/path/cosmos-embed/msrvtt-subset/
├── msrvtt_test_1k.json
└── video/
    ├── video7020.mp4
    └── ...

将该前缀下载/同步到挂载的 data/ 目录后，使用上文相同的 Docker 命令。

输出

results/
├── train/
│   ├── cosmos_embed1_model_latest.pth
│   ├── cosmos_embed1_model_<iter>.pth
│   └── experiment.yaml
├── evaluate/
│   ├── metrics.json
│   └── experiment.yaml
├── inference/
│   ├── results.json
│   └── experiment.yaml
├── export/
│   ├── cosmos_embed1_combined.onnx
│   └── export_config.yaml
└── export_hf/
    └── cosmos_embed1_hf/

已知问题

现象	原因	修复方法
MSRVTTDataset: 0 videos found	mp4_urls 不是本地 glob，或元数据文件名与视频不匹配。	将数据挂载到容器中，并设置 mp4_urls=/data/video/*.mp4。
HF 下载/身份验证失败	缺少或无效的 HF_TOKEN，或尚未接受模型协议。	接受模型条款，并传入 -e HF_TOKEN。
LoRA 注入失败	Transformer Engine 视觉编码器已启用。	设置 model.network.visual_encoder.transformer_engine=false。
ONNX/HF 导出提示缺少组件	导出检查点不完整或仅包含适配器。	使用完整检查点，或在导出前配置预训练视觉/文本来源。
CUDA OOM	批次/分辨率对于 GPU 来说过高。	减小批次大小，使用 224p，启用 LoRA，或使用更多 GPU。