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 にあります。

Train アクションのポリシー

schemas/ に Cosmos-Embed のスキーマがないため、このモデルスキルには AutoML がパッケージ化されていません。上位レベルのリクエストに automl_policy: on が含まれていても、train、evaluate、inference、export には常にモデルスキルの直接アクションを使用してください。モデル固有の train スキーマとテンプレートが追加されるまでは、Cosmos-Embed をワークフローまたは AutoML スキル経由で実行しないでください。

evaluate、inference、export、デプロイフローなどの train 以外のアクションは、このモデルスキルにとどまります。実行ごとの 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 イメージは、Hugging Face キャッシュの変数で別の場所が指定されていても、BERT Q-Former コンポーネントを /usr/local/lib/python3.12/dist-packages/pretrained_checkpoints にダウンロードします。上記の専用バインドマウントは、プラットフォームスキルで選択されたコンテナー 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"

イメージには protobuf==7.x を使用する wandb==0.21.0 が含まれており、protobuf を 7 未満に固定しないと、トレーニング／評価の前に W&B のインポートが失敗します。スモークテストまたはオフライン実行では、WANDB_DISABLED=true と WANDB_MODE=disabled を使用してください。モデルのチェックポイントが無効でも、Cosmos-Embed は公開されている google-bert/bert-base-uncased の Q-Former コンポーネントをダウンロードする場合があります。そのため、環境変数として HF_TOKEN を渡すか、永続的な Hugging Face キャッシュをマウントしてください。トークンを仕様、ログ、レポートに書き込まないでください。

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"

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"

Hugging Face 形式でエクスポート：

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 を取得し、メタデータに存在する動画だけを対象にします。実行時に動画が 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/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 ファイルを解決して渡してください。アクション仕様のテンプレートでは、これらのチェックポイントフィールドは意図的に null のままにしてあるため、モデルスキルランナーまたはユーザーがリゾルバーで選択したチェックポイントを指定する必要があります。latest を明示的に求められた場合に限り、最新のシンボリックリンクを使用してください。

統合済みチェックポイントから単一 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	単一 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 を使用する実行では、まずデータセットの一部または全体を実行ホスト／コンテナーのファイルシステムにステージングしてから、仕様には /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 の OOM	バッチサイズまたは解像度が GPU に対して大きすぎます。	バッチサイズを減らす、224p を使用する、LoRA を有効にする、または GPU を増やしてください。