nvidia-tao/tao-skill-bank/tao-finetune-cosmos-embed

Cosmos-Embed

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

Cosmos-Embed1 is a joint video-text embedder for text-to-video retrieval, video-to-video search, zero-shot/kNN classification, and semantic deduplication. The packaged CLI is cosmos-embed1 and supports train, evaluate, inference, and export.

Container image and per-action commands are in references/skill_info.yaml. Compact starting specs are in references/spec_template_*.yaml; the LoRA fine-tuning recipe on chunk manifests is in references/lora-finetune-physicalai.md with references/spec_template_train_lora.yaml.

Release note: the pinned 7.3.0-cosmos-embed image is the first that trains ViT LoRA adapters, accepts caption_mode, label_mode, and train.grad_accum_iter, scores the in-training retrieval_benchmark, and validates large ONNX exports by file path. On 7.1.0-cosmos-embed the ViT adapters silently stay at zero and the vision ONNX export exits 1 after writing valid files.

Train Action Policy

AutoML is not packaged for this model skill (no Cosmos-Embed schemas under schemas/). Use the direct train, evaluate, inference, and export actions even when a higher-level request includes automl_policy: on; the per-run override does not change model metadata, and non-train and deploy flows stay in this skill.

Quick Start

Use the published container pinned below (also declared in references/skill_info.yaml); do not build from the Cosmos-Embed1 source tree for normal skill use.

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"

Expected local workspace layout:

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

Use these Docker options for all actions unless the local Docker/platform skill gives a stricter environment-specific command:

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 --ipc=host --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"
)

The image downloads the BERT Q-Former component under /usr/local/lib/python3.12/dist-packages/pretrained_checkpoints regardless of the Hugging Face cache variables, so the bind mount above must exist and be writable by the container UID chosen by the platform skill. Verify the cache location before reusing this path with another image version.

Keep -w /results: with a cold Triton cache, a working directory containing specs/ makes GCC read ./specs as a compiler spec file (cannot read spec file './specs': Is a directory).

For Cosmos-Embed images that ship protobuf==7.x, run a small startup preamble before every action:

python -m pip install "protobuf<7"

wandb==0.21.0 cannot import with protobuf==7.x; the preamble is a no-op where protobuf is already below 7. Keep WANDB_DISABLED=true for smoke or offline runs. The Q-Former path may download the public google-bert/bert-base-uncased even with the model checkpoint disabled, so provide HF_TOKEN through the environment or an --env-file, or mount a persistent HuggingFace cache. Never write the token into specs, logs, or reports.

Train:

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"

Evaluate:

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"

Inference:

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"

Export ONNX (one graph per tower, dynamic batch axis; run once with export.mode=video and once with export.mode=text):

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.mode=video export.batch_size=-1 export.on_cpu=true \
  export.onnx_file=/results/export/cosmos_embed1_vision.onnx \
  results_dir=/results"

Export HuggingFace format:

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"

For a LoRA checkpoint, every downstream spec (evaluate, inference, export) must carry the training spec's model.lora block and list exactly the modules that were trained in target_modules; export merges the adapters into the weights. See references/lora-finetune-physicalai.md.

Smoke Overrides

For a quick functional check of any action, keep the same specs and shrink iterations, batch sizes, and workers with the command-line overrides listed in references/smoke-overrides.md (including the scheduler and no-checkpoint caveats).

Data Format

The MSR-VTT path expects a local video glob and a JSON metadata file:

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

List-format metadata rows must include at least video and caption:

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

The dataset loader derives the video id from the local .mp4 filename and filters to videos present in the metadata. If a run finds zero videos, check that mp4_urls points to a container-local glob and that metadata video names match the filenames.

Chunk manifests (dataset_type: vad_r1_chunks, the PhysicalAI-Event-Videos labels.json layout) hold one record per labeled chunk of a parent video, with caption fields that may be lists of strings. caption_field names the fields that feed the caption pool, caption_mode turns the pool into samples, label_mode decides what counts as a negative, and path_prefix_mapping rewrites the manifest's video prefix to the mounted folder. dataset_type: retrieval_benchmark (validation and evaluation only) takes a benchmark release root holding gt/, additional_gt/, clips/. Details and staging layouts: references/lora-finetune-physicalai.md, references/data-staging.md.

Model Weights

• Local HF directory: mount it under /model and set model.pretrained_model_path=/model/Cosmos-Embed1-224p.
• HuggingFace repo: set model.pretrained_model_path=nvidia/Cosmos-Embed1-224p (or -336p, -448p) and pass HF_TOKEN if access is gated.
• Fine-tuned checkpoint: set downstream actions to the resolver-selected /results/train/checkpoints/iter_#########.pt file.

Training writes full checkpoints under results/train/checkpoints/iter_#########.pt, updates results/train/checkpoints/latest_checkpoint.txt, and creates a cosmos_embed1_model_latest.pth symlink. For evaluate.checkpoint, inference.checkpoint, export.checkpoint, and train.resume_training_checkpoint_path, resolve and pass the exact iter_#########.pt file for the intended iteration. The action spec templates intentionally leave these checkpoint fields null so the model-skill runner or the user must provide the resolver-selected checkpoint. Use the latest symlink only when the user explicitly asks for latest. When the retrieval_benchmark callback is enabled, results/train/selection.json names the best-scoring checkpoint; prefer it over the latest one.

For single-GPU resume/retrain from a consolidated checkpoint, set model.fsdp_shard_size: 1. The container default is 8, which sends resumed training through an FSDP apply path that Cosmos-Embed1 does not implement for this model class. Do not resume a LoRA run through train.resume_training_checkpoint_path; export the selected checkpoint to HF format, point model.pretrained_model_path at it, and start a fresh LoRA stage with its own schedule.

Variants:

Keep model.network.embed_dim, model.input_hw, and model.network.spatial_resolution aligned with the selected variant.

Important Parameters