tao-finetune-cosmos-embed

Cosmos-Embed

单独安装？ 如果此会话并非由 TAO 技能库插件初始化，请先运行 tao-setup 技能（主机预检、凭据、跨技能发现）。

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

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

训练操作策略

由于 schemas/ 下没有 Cosmos-Embed 的架构，因此此模型技能未打包 AutoML。始终直接使用模型技能操作执行 train、evaluate、inference 和 export，即使更高层级的请求包含 automl_policy: on 也是如此。在添加特定于模型的训练架构和模板之前，请勿通过工作流或 AutoML 技能运行 Cosmos-Embed。

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

快速开始

使用下方固定版本的已发布 Cosmos-Embed 容器（也在 references/skill_info.yaml 中声明）。常规使用此技能时，请勿从私有的 Cosmos-Embed1 源代码树构建；仅在开发容器本身时才从源代码构建。

COSMOS_EMBED_IMAGE_DEFAULT=nvcr.io/nvidia/tao/tao-toolkit:7.3.0-cosmos-embed  # versions-key: images.tao_toolkit.cosmos_embed
COSMOS_EMBED_IMAGE="${COSMOS_EMBED_IMAGE:-$COSMOS_EMBED_IMAGE_DEFAULT}"
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 选项：

set -a; source /path/to/.env; set +a   # omit if already exported
COSMOS_EMBED_IMAGE_DEFAULT=nvcr.io/nvidia/tao/tao-toolkit:7.3.0-cosmos-embed  # versions-key: images.tao_toolkit.cosmos_embed
COSMOS_EMBED_IMAGE="${COSMOS_EMBED_IMAGE:-$COSMOS_EMBED_IMAGE_DEFAULT}"
RUN_ROOT="${RUN_ROOT:-$PWD}"
COSMOS_EMBED_PRETRAINED_CACHE="${COSMOS_EMBED_PRETRAINED_CACHE:-$RUN_ROOT/pretrained_checkpoints}"
mkdir -p "$COSMOS_EMBED_PRETRAINED_CACHE"
DOCKER_COMMON=(
  --rm --gpus all --shm-size=8g --network=host
  --shm-size=64g
  --ulimit memlock=-1
  --ulimit stack=67108864
  -w /results
  -e HF_TOKEN
  -e WANDB_DISABLED=true
  -e WANDB_MODE=disabled
  -e HUGGINGFACE_HUB_CACHE=/hf_cache
  -v "$RUN_ROOT/data:/data:ro"
  -v "$RUN_ROOT/model:/model"
  -v "$RUN_ROOT/specs:/specs:ro"
  -v "$RUN_ROOT/results:/results"
  -v "$RUN_ROOT/hf_cache:/hf_cache"
  --mount "type=bind,source=$COSMOS_EMBED_PRETRAINED_CACHE,target=/usr/local/lib/python3.12/dist-packages/pretrained_checkpoints"
)

固定版本的 7.1.0 镜像会将 BERT Q-Former 组件下载到 /usr/local/lib/python3.12/dist-packages/pretrained_checkpoints，即使 Hugging Face 缓存变量指向其他位置也是如此。上方专用绑定挂载目录必须可由 平台技能选定的容器 UID 写入；请在提交前以启动用户身份创建该目录。应用 平台的非 root UID/GID 和可写 HOME 设置时，请保留此挂载。这样即可支持正常的 运行时下载，而无需修改镜像或使整个 site-packages 目录可写。使用其他镜像版本时，请先确认其缓存位置，再复用此镜像专用的目标路径。

使用可写的每次运行结果目录作为容器工作目录。此镜像的 Triton 缓存为空时， 如果从包含 specs/ 的目录启动，GCC 可能会将 ./specs 读取为编译器规范文件并报错： cannot read spec file './specs': Is a directory。请保持推理配置中的绝对 results_dir 不变；仅更改工作目录即可。

对于附带 protobuf==7.x 的 Cosmos-Embed 镜像，请在每项操作之前运行一段简短的启动预处理命令：

python -m pip install "protobuf<7"

镜像包含 wandb==0.21.0 和 protobuf==7.x；除非将 protobuf 固定为 7 以下版本，否则导入 W&B 会在训练/评估前失败。进行冒烟测试或离线运行时，请使用 WANDB_DISABLED=true 和 WANDB_MODE=disabled。即使模型检查点已禁用，Cosmos-Embed 仍可能下载公开的 google-bert/bert-base-uncased Q-Former 组件，因此请将 HF_TOKEN 作为环境变量传入，或挂载持久化的 HuggingFace 缓存。请勿将令牌写入配置、日志或报告。

训练：

docker run "${DOCKER_COMMON[@]}" "$COSMOS_EMBED_IMAGE" \
  bash -lc "python -m pip install 'protobuf<7' && cosmos-embed1 train -e /specs/train.yaml results_dir=/results"

评估：

docker run "${DOCKER_COMMON[@]}" "$COSMOS_EMBED_IMAGE" \
  bash -lc "python -m pip install 'protobuf<7' && cosmos-embed1 evaluate -e /specs/evaluate.yaml results_dir=/results"

推理：

docker run "${DOCKER_COMMON[@]}" "$COSMOS_EMBED_IMAGE" \
  bash -lc "python -m pip install 'protobuf<7' && 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" \
  bash -lc "python -m pip install 'protobuf<7' && cosmos-embed1 export -e /specs/export_onnx.yaml \
  export.checkpoint=/results/train/checkpoints/iter_000000001.pt \
  export.onnx_file=/results/export/cosmos_embed1_combined.onnx \
  results_dir=/results"

导出 HuggingFace 格式：

docker run "${DOCKER_COMMON[@]}" "$COSMOS_EMBED_IMAGE" \
  bash -lc "python -m pip install 'protobuf<7' && cosmos-embed1 export -e /specs/export_hf.yaml \
  export.checkpoint=/results/train/checkpoints/iter_000000001.pt \
  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
train.optim.warmup_steps=0
train.optim.lr_decay_iters=1
dataset.train_dataset.batch_size=1
dataset.val_dataset.batch_size=1
dataset.train_dataset.workers=0
dataset.val_dataset.workers=0

为冒烟测试缩短余弦调度器时，确保 train.optim.lr_decay_iters 大于 train.optim.warmup_steps，或像上面那样设置 train.optim.warmup_steps=0。调度器会除以 lr_decay_iters - warmup_steps，因此两者相等时会在检查点写入前失败。

如果没有本地 Cosmos-Embed1 预训练检查点，可将 model.pretrained_model_path=null，进行仅检查流程连通性的冒烟训练。在这种模式下，模型质量没有意义，但仍可运行训练、评估、推理和导出操作路径。在当前容器中，Q-Former 路径仍可获取 google-bert/bert-base-uncased；请为全新的临时容器提供 HF_TOKEN 或挂载 HuggingFace 缓存。

针对小型子集进行评估和推理冒烟测试：

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/checkpoints/iter_#########.pt 文件。

训练会将完整检查点写入 results/train/checkpoints/iter_#########.pt，更新 results/train/checkpoints/latest_checkpoint.txt，并创建 cosmos_embed1_model_latest.pth 符号链接。对于 evaluate.checkpoint、inference.checkpoint、export.checkpoint 和 train.resume_training_checkpoint_path，请解析并传入目标迭代对应的确切 iter_#########.pt 文件。操作配置模板会有意将这些检查点字段留空，因此模型技能运行器或用户必须提供解析器选定的检查点。只有在用户明确要求使用最新版本时，才使用最新版本符号链接。

从整合后的检查点进行单 GPU 恢复训练或重新训练时，请设置 model.fsdp_shard_size: 1。容器默认值为 8，这会使恢复训练走入 Cosmos-Embed1 尚未为此模型类实现的 FSDP 应用路径。

变体：

变体	分辨率	帧数	嵌入维度
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	1 表示使用单块 GPU，>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 开始；除非实验结果另有显示，否则将其保持在接近 2 * lora_rank 的值。
model.lora.lora_dropout	LoRA dropout。从 0.1 开始；对于小型数据集，可搜索 0.0、0.05 和 0.1。
model.lora.bias	Bias 策略：none、all 或 lora_only。除非有意训练 bias，否则保持为 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 的运行，请先将部分或全部数据集暂存到执行主机/容器的文件系统中，然后在配置中使用 /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。
cannot import name 'Imports' from 'wandb.proto.wandb_telemetry_pb2'	容器中的 wandb==0.21.0 与 protobuf==7.x 不兼容。	在容器中调用 cosmos-embed1 之前运行 python -m pip install "protobuf<7"。
使用 Model does not implement 'apply_fsdp' 恢复失败	单 GPU 恢复时加载了合并后的检查点，而 model.fsdp_shard_size 仍为默认值 8。	对本地单 GPU 恢复/重新训练，将 model.fsdp_shard_size=1。
LoRA 注入失败	已启用 Transformer Engine 视觉编码器。	设置 model.network.visual_encoder.transformer_engine=false。
ONNX/HF 导出提示缺少组件	导出检查点不完整或仅包含适配器。	使用完整检查点，或在导出前配置预训练视觉/文本源。
CUDA 内存不足	对 GPU 而言，批次/分辨率过高。	减小批次大小、使用 224p、启用 LoRA 或增加 GPU 数量。