Fine-tune Cosmos pour générer des vidéos de robots

Pourquoi tu as besoin de ça

Cosmos Predict 2.5 est un modèle qui génère des vidéos à partir de texte, d'images ou de clips vidéo. Il marche bien sur du contenu générique, mais si tu veux l'utiliser pour ta gamme de robots — disons un bras manipulateur dans un atelier spécifique — il faut l'adapter.

Le problème : fine-tuner 2 milliards de paramètres coûte une fortune en GPU et en mémoire. Pire, tu risques d'effacer tout ce qu'il savait du monde réel (on appelle ça le « catastrophic forgetting »).

Comment LoRA et DoRA changent la donne

LoRA (Low-Rank Adaptation) ajoute de petits modules d'adaptation qu'on entraîne, tandis que le modèle de base reste figé. C'est comme brancher un adaptateur sur un vieux lecteur VHS : le lecteur ne change pas, mais tu peux lire un nouveau format.

DoRA c'est une variante : elle décompose chaque poids en magnitude et direction avant d'appliquer la mise à jour rang faible. En pratique, ça converge souvent plus vite.

Résultat : tu peux fine-tuner sur une H100 seule (voire moins si tu fais du gradient checkpointing), et tes fichiers d'adaptation tiennent en quelques Mo — tu peux en changer à l'inférence selon le robot ou la tâche.

Le setup

Tu vas avoir besoin :

• ▸Python 3.10+, PyTorch 2.5+ avec CUDA
• ▸diffusers, transformers, peft, accelerate
• ▸Minimum 80 Go de VRAM (un GPU H100) ; 8× H100 c'est mieux mais pas obligatoire

pip install -U "diffusers[torch]" transformers accelerate peft wandb

NVIDIA fournit un exemple concret : 92 vidéos de tâches de saisie et dépôt d'objets avec descriptions textuelles, plus 50 paires (image + prompt) pour l'évaluation. C'est petit, c'est intentionnel — LoRA fonctionne avec peu de données.

Comment ça marche en pratique

1. ▸

   Charger le modèle et geler les poids principaux

   • ▸Le VAE (encode les vidéos en latents), l'encodeur de texte, le DiT (diffusion transformer) restent figés
   • ▸Seul le DiT reçoit les modules LoRA/DoRA sur ses couches d'attention et feedforward
2. ▸

   Entraîner avec rectified flow

   • ▸Cosmos utilise une technique appelée rectified flow : on bruite les latents et le modèle apprend à prédire la vélocité qui ramène le bruit vers la vidéo propre
   • ▸Les deux premières images servent de conditions (pas d'ajout de bruit), ça aide le modèle à bien ancrer la génération
3. ▸

   Sauvegarder et permuter

   • ▸Tous les 100 epochs, le script sauvegarde les poids LoRA (format diffusers)
   • ▸À l'inférence, tu charges le module d'adaptation que tu veux — un pour les bras, un pour les pinces, etc.

Lancer l'entraînement

export MODEL_NAME="nvidia/Cosmos-Predict2.5-2B"
export DATA_DIR="gr1_dataset/train"
export OUT_DIR=YOUR_OUTPUT_DIR

accelerates launch --mixed_precision="bf16" train_cosmos_predict25_lora.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --train_data_dir=$DATA_DIR \
  --train_batch_size=1 \
  --num_train_epochs=500 \
  --lora_rank 32

Le batch size est 1 parce que les vidéos sont lourdes en mémoire. Le rang LoRA à 32 c'est un bon compromis : assez expressif, assez petit pour rester efficace.

Après l'entraînement

Tu utilises le modèle fine-tuné pour générer des vidéos synthétiques de trajectoires robots. Pas besoin d'attendre qu'un vrai robot manipule 1000 fois le même objet — tes données synthétiques entraînent tes algorithmes de contrôle.

C'est particulièrement utile si tu veux explorer des scénarios rares, des erreurs, des cas limites : tu décris « le bras doit attraper l'objet même s'il est légèrement tourné » et le modèle génère des variantes pour toi.