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.

Train Action Policy

AutoML is not packaged for this model skill because there are no Cosmos-Embed schemas under schemas/. Always use the direct model skill actions for train, evaluate, inference, and export, even when a higher-level request includes automl_policy: on. Do not route Cosmos-Embed through workflow or AutoML skills until model-specific train schemas and templates are added.

Non-train actions such as evaluate, inference, export, and deploy flows stay in this model skill. The per-run automl_policy override does not change model metadata.

Quick Start

Use the published Cosmos-Embed container pinned below (also declared in references/skill_info.yaml). Do not build from the private Cosmos-Embed1 source tree for normal skill use; build from source only when developing the container itself.

COSMOS_EMBED_IMAGE_DEFAULT=nvcr.io/nvidia/tao/tao-toolkit:7.1.0-cosmos-embed  # versions-key: images.tao_toolkit.cosmos_embed
COSMOS_EMBED_IMAGE="${COSMOS_EMBED_IMAGE:-$COSMOS_EMBED_IMAGE_DEFAULT}"
docker pull

Expected local workspace layout:

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

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

The image contains wandb==0.21.0 with protobuf==7.x; importing W&B fails before training/evaluation unless protobuf is pinned below 7. Use WANDB_DISABLED=true and WANDB_MODE=disabled for smoke or offline runs. Cosmos-Embed may still download the public google-bert/bert-base-uncased Q-Former component even when the model checkpoint is disabled, so pass HF_TOKEN as an environment variable or mount a persistent HuggingFace cache. Do not write the token into specs, logs, or reports.

Train:

Evaluate:

Inference:

Export ONNX:

Export HuggingFace format:

Smoke Overrides

For a small functional check, keep the same specs and override the expensive knobs:

When shortening the cosine scheduler for smoke runs, keep train.optim.lr_decay_iters greater than train.optim.warmup_steps, or set train.optim.warmup_steps=0 as shown above. The scheduler divides by lr_decay_iters - warmup_steps, so equal values fail before the checkpoint is written.

If no local Cosmos-Embed1 pretrained checkpoint is available, set model.pretrained_model_path=null for a plumbing-only smoke train. The model quality is meaningless in that mode, but the train/evaluate/inference/export action paths can still be exercised. In the current container, the Q-Former path can still fetch google-bert/bert-base-uncased; provide HF_TOKEN or a mounted HuggingFace cache for fresh ephemeral containers.

For evaluation and inference smoke tests on a tiny subset:

Data Format

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

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

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.

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 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.

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.

Variants:

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

Important Parameters

LoRA and AutoML Notes

For parameter-efficient fine-tuning, set model.lora.enabled=true and keep model.network.visual_encoder.transformer_engine=false; TAO Core's Cosmos-Embed1 config notes that PEFT cannot inject adapters into Transformer Engine layers. Treat the LoRA fields above as the first candidate parameters for manual tuning or AutoML-style search before unfreezing larger model blocks. Avoid changing target_modules or modules_to_save unless the user explicitly needs custom adapter placement.

S3 Staging

The Cosmos-Embed1 CLI consumes local paths and Python globs, not raw s3://.../*.mp4 URIs. For S3-backed runs, first stage a subset or full dataset to the execution host/container filesystem, then use local paths such as /data/video/*.mp4 in the spec.

Recommended S3 layout for staged MSR-VTT data:

After downloading/syncing that prefix into the mounted data/ directory, use the same Docker commands above.

Outputs

Known Pitfalls