Cosmos-Embed

Cosmos-Embed1 est un embedder vidéo-texte joint pour la récupération texte-vers-vidéo, la recherche vidéo-vers-vidéo, la classification zero-shot/kNN et la déduplication sémantique. Le CLI packagé est cosmos-embed1 et supporte train, evaluate, inference et export.

L'image conteneur et les commandes par action sont dans references/skill_info.yaml. Les specs de démarrage compact sont dans references/spec_template_*.yaml.

Train Action Policy

Ce modèle est AutoML-enabled au niveau du modèle. Avant de traiter toute demande au stade train, lisez references/skill_info.yaml et résolvez l'override de run depuis soit une valeur automl_policy explicite, soit la demande de workflow de l'utilisateur. Traitez les phrases comme "turn off AutoML", "disable AutoML", "no HPO", ou "plain training" comme automl_policy: off pour cette run uniquement ; sinon, par défaut auto. Quand automl_policy: auto, automl_enabled: true, et que schemas/train.schema.json et references/spec_template_train.yaml sont packagés, routez l'action train via tao-skill-bank:tao-run-automl par défaut avec le skill_dir de ce modèle. Préservez les overrides de workflow/application pour les datasets, specs, répertoires de sortie, paramètres GPU/platform, checkpoints parents et automl_policy. Utilisez l'entraînement direct du modèle uniquement quand automl_policy: off ou que le schema/template train packagé manque ; en cas de schema manquant, signalez que AutoML est activé mais non exécutable pour ce modèle jusqu'à ce que les schemas soient générés.

Les actions non-train comme evaluate, inference, export et les flows de déploiement restent dans cette skill de modèle. L'override automl_policy par run ne change pas les métadonnées du modèle.

Quick Start

Utilisez le conteneur Cosmos-Embed publié déclaré par references/skill_info.yaml et résolu via versions.yaml. Ne faites pas de build à partir de l'arborescence source privée de Cosmos-Embed1 pour un usage normal de skill ; faites un build à partir de la source uniquement lors du développement du conteneur lui-même.

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"

Layout du workspace local attendu :

workspace/
├── data/
│   ├── msrvtt_test_1k.json
│   └── video/
│       ├── video7020.mp4
│       └── ...
├── model/
│   └── Cosmos-Embed1-224p/        # optionnel si on utilise l'id repo HF
├── specs/
│   ├── train.yaml
│   ├── evaluate.yaml
│   ├── inference.yaml
│   ├── export_onnx.yaml
│   └── export_hf.yaml
└── results/

Utilisez ces options Docker pour toutes les actions sauf si la skill Docker/platform locale donne une commande plus stricte spécifique à l'environnement :

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"
)

Train :

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

Evaluate :

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

Inference :

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

Export 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

Export format 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

Smoke Overrides

Pour un petit contrôle fonctionnel, gardez les mêmes specs et overridez les paramètres coûteux :

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

Si aucun checkpoint pré-entraîné local Cosmos-Embed1 ou token HuggingFace n'est disponible, définissez model.pretrained_model_path=null pour un smoke train de plomberie uniquement. La qualité du modèle est insignifiante dans ce mode, mais les chemins d'action train/evaluate/inference/export peuvent toujours être exercés.

Pour les smoke tests d'evaluate et inference sur un petit sous-ensemble :

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

Data Format

Le chemin MSR-VTT attend un glob vidéo local et un fichier de métadonnées JSON :

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

Les lignes de métadonnées au format liste doivent inclure au minimum video et caption :

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

Le dataset loader déduit l'id vidéo du nom de fichier .mp4 local et filtre les vidéos présentes dans les métadonnées. Si une run trouve zéro vidéo, vérifiez que mp4_urls pointe vers un glob local au conteneur et que les noms video des métadonnées correspondent aux noms de fichiers.

Model Weights

• Répertoire HF local : montez-le sous /model et définissez model.pretrained_model_path=/model/Cosmos-Embed1-224p.
• Repo HuggingFace : définissez model.pretrained_model_path=nvidia/Cosmos-Embed1-224p et passez HF_TOKEN si l'accès est gardienné.
• Checkpoint fine-tuné : les actions en aval par défaut à /results/train/cosmos_embed1_model_latest.pth.

Variantes :

Gardez model.network.embed_dim, model.input_hw et model.network.spatial_resolution alignés avec la variante sélectionnée.

Important Parameters

LoRA and AutoML Notes

Pour le fine-tuning parameter-efficient, définissez model.lora.enabled=true et gardez model.network.visual_encoder.transformer_engine=false ; la config TAO Core de Cosmos-Embed1 note que PEFT ne peut pas injecter d'adaptateurs dans les couches Transformer Engine. Traitez les champs LoRA ci-dessus comme les premiers paramètres candidats pour le tuning manuel ou la recherche de style AutoML avant de dégeler des blocs de modèle plus larges. Évitez de changer target_modules ou modules_to_save sauf si l'utilisateur a explicitement besoin d'un placement d'adaptateur personnalisé.

S3 Staging

Le CLI Cosmos-Embed1 consomme des chemins locaux et des globs Python, pas des URIs bruts s3://.../*.mp4. Pour les runs avec S3, staging d'abord un sous-ensemble ou le dataset complet sur le filesystem hôte/conteneur d'exécution, puis utilisez des chemins locaux comme /data/video/*.mp4 dans la spec.

Layout S3 recommandé pour les données MSR-VTT stagées :

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

Après téléchargement/sync de ce préfixe dans le répertoire data/ monté, utilisez les mêmes commandes Docker ci-dessus.

Outputs

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/

Known Pitfalls