tao-finetune-cosmos-embed

Cosmos-Embed

Cosmos-Embed1 は、テキストから動画への検索、動画間検索、ゼロショット／kNN 分類、セマンティック重複排除に対応する、動画とテキストの共同埋め込みモデルです。パッケージ化された CLI は cosmos-embed1 で、train、evaluate、inference、export をサポートしています。

コンテナイメージとアクションごとのコマンドは references/skill_info.yaml にあります。コンパクトな開始用仕様は references/spec_template_*.yaml にあります。

Train アクションポリシー

このモデルはモデルレイヤーで AutoML が有効になっています。train ステージのリクエストに対応する前に、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 を指定して、train アクションを tao-skill-bank:tao-run-automl 経由で実行してください。データセット、仕様、出力ディレクトリ、GPU／プラットフォーム設定、親チェックポイント、automl_policy に関するワークフロー／アプリケーションのオーバーライドは維持してください。automl_policy: off の場合、またはパッケージ化された train スキーマ／テンプレートが存在しない場合に限り、モデルの直接トレーニングを使用してください。スキーマがない場合は、AutoML は有効ですが、スキーマが生成されるまではこのモデルで実行できないことを報告してください。

evaluate、inference、export、デプロイフローなどの train 以外のアクションは、このモデルのスキル内で実行します。実行ごとの 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"
)

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

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 の各アクション経路は引き続き確認できます。

小さなサブセットを使った evaluation と inference のスモークテスト：

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 を導出し、メタデータに存在する動画に絞り込みます。実行時に動画が 0 件見つかった場合は、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 から始め、実験で別の値が適切と判明しない限り、2 * lora_rank に近い値を保ってください。
model.lora.lora_dropout	LoRA のドロップアウトです。まずは 0.1 から始め、小規模なデータセットでは 0.0、0.05、0.1 を試してください。
model.lora.bias	バイアスのポリシーです。none、all、lora_only を指定できます。意図的にバイアスを学習させる場合を除き、none のままにしてください。
model.lora.use_rslora / use_dora	オプションの LoRA バリアントです。一度に 1 つだけ有効にし、チェックポイントとともに設定を記録してください。
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	デバイスの不一致に関する問題を回避するため、エクスポート時に推奨されます。

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 を増やしてください。