NVIDIA Cosmos Predict 2.5 : LoRA et DoRA pour la robotique

Introduction à NVIDIA Cosmos Predict 2.5

NVIDIA Cosmos Predict 2.5 est un modèle de monde à grande échelle, conçu pour générer des vidéos qui sont physiquement plausibles. Ces vidéos peuvent être conditionnées par du texte, des images ou même des clips vidéo. Cependant, pour adapter ce modèle à des domaines spécifiques, comme la manipulation robotique ou un point de vue de caméra particulier, il est nécessaire de procéder à un ajustement ciblé.

L'entraînement de politiques robotiques nécessite des données de démonstration. Or, la collecte de trajectoires de robots réels est un processus à la fois lent et coûteux. Une alternative consiste à générer des trajectoires synthétiques en utilisant un modèle vidéo de monde ajusté. Toutefois, ajuster complètement un modèle qui contient 2 milliards de paramètres est une tâche coûteuse. De plus, cela peut entraîner un oubli catastrophique des connaissances générales acquises par le modèle.

LoRA et DoRA : Des solutions innovantes

Pour résoudre ce problème, les méthodes LoRA et DoRA ont été développées. Ces méthodes permettent d'injecter de petits modules adaptateurs entraînables dans le modèle de base, qui est gelé. Cela réduit les exigences en mémoire tout en maintenant les fichiers d'adaptateurs petits et portables. Grâce à cette approche, l'ajustement devient pratique sur un seul GPU, et il est possible d'échanger facilement les adaptateurs pour différents domaines lors de l'inférence.

Le guide explique comment procéder à un ajustement efficace en paramètres de Cosmos Predict 2.5 avec LoRA et DoRA. Pour cela, il utilise les bibliothèques diffusers et accelerate, qui supportent l'entraînement sur un ou plusieurs GPU. Après cet ajustement, le modèle peut être utilisé pour générer des trajectoires robotiques synthétiques, utiles pour des tâches d'apprentissage en robotique en aval.

Prérequis techniques

Pour mettre en œuvre cet ajustement, il est nécessaire d'avoir PyTorch 2.5+ avec CUDA, ainsi que les bibliothèques diffusers et accelerate. L'installation de wandb est optionnelle mais recommandée pour surveiller l'entraînement. Au minimum, un GPU de 80 Go est requis pour l'entraînement sur un seul GPU. Cependant, pour une itération plus rapide, 8 GPU H100 sont recommandés.

Voici comment installer les dépendances sur votre machine :

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

Après avoir installé diffusers, il est conseillé de naviguer vers examples/cosmos pour explorer le code d'exemple. Les ensembles de données utilisés pour l'entraînement et le test comprennent 92 vidéos de manipulation robotique avec des invites textuelles décrivant des tâches de prise et de placement, ainsi que 50 paires (invite, image) pour le test.

Pour télécharger et prétraiter les ensembles de données d'entraînement et de test, utilisez le script suivant :

bash download_and_preprocess_datasets.sh

Le dossier résultant de l'ensemble de données d'entraînement ressemble à ceci :

gr1_dataset/train
└── metadata.csv

L'ensemble de données d'évaluation est un répertoire plat de fichiers .txt et .png appariés pour les paires (invite, image) :

gr1_dataset/test
├── filename1.txt
├── filename1.png
├── filename2.txt
├── filename2.png

Mise en œuvre et entraînement

Dans cette section, nous parcourons l'implémentation dans train_cosmos_predict25_lora.py. VideoDataset charge chaque échantillon comme une paire (légende, vidéo) depuis args.train_data_dir (gr1_dataset/train dans notre exemple). Pour les vidéos plus longues que args.num_frames, elle échantillonne une fenêtre continue aléatoire de args.num_frames à chaque époque, permettant une augmentation temporelle. En interne, VideoProcessor de diffusers.video_processor redimensionne et normalise les images brutes en un tenseur de forme (canaux, images, hauteur, largeur).

train_dataset = VideoDataset(
    dataset_dir=args.train_data_dir,
    num_frames=args.num_frames,
    video_size=[args.height, args.width],
)

Cosmos Predict 2.5 se compose de trois sous-modules :

• Un VAE qui encode les vidéos en latents
• Un encodeur de texte qui encode les invites textuelles en embeddings d'invite
• DiT pour la diffusion dans l'espace latent

Pendant l'entraînement, tous les poids du VAE, de l'encodeur de texte et de DiT sont gelés. Les adaptateurs LoRA sont injectés dans les projections d'attention de DiT (to_q, to_k, to_v, to_out.0) et dans les couches feedforward (ff.net.0.proj, ff.net.2). Les paramètres LoRA entraînables sont ensuite convertis en float32 pour la stabilité numérique sous précision mixte bf16.

from diffusers import Cosmos2_5_PredictBasePipeline
from peft import LoraConfig

pipe = Cosmos2_5_PredictBasePipeline.from_pretrained(
    "nvidia/Cosmos-Predict2.5-2B",
    revision="diffusers/base/post-trained",
    torch_dtype=torch.bfloat16,
)

dit = pipe.[transformer](/glossaire/transformer)
text_encoder = pipe.text_encoder
dit.requires_grad_(False)
vae.requires_grad_(False)
text_encoder.requires_grad_(False)

lora_config = LoraConfig(
    r=args.lora_rank,
    lora_alpha=args.lora_alpha,
    target_modules=['to_q', 'to_k', 'to_v', 'to_out.0', 'ff.net.0.proj', 'ff.net.2'],
    use_dora=args.use_dora,
)

dit.add_adapter(lora_config) cast_training_params(dit, dtype=torch.float32)

Passer **use_dora=True** bascule vers DoRA, qui décompose chaque poids en magnitude et direction avant d'appliquer la mise à jour de faible rang. Aucune autre modification de la boucle d'entraînement n'est nécessaire.

**Cosmos Predict 2.5** utilise un **flux rectifié** : le modèle est entraîné pour prédire la vitesse qui transporte linéairement un échantillon de bruit vers les données "propres" d'origine. Concrètement, au moment t, une interpolation bruyante **xt** est construite à un niveau de bruit échantillonné **σt**, et le modèle apprend à prédire la vitesse cible **bruit - propre** via les erreurs quadratiques moyennes (MSE loss). Les deux premières images de la vidéo sont utilisées comme conditionnement, et donc aucun bruit n'est ajouté à leurs latents.

La perte d'entraînement suit la formulation de flux rectifié utilisée par **Cosmos Predict 2.5** :

```python
sigma_t = sample_train_sigma_t(bsz, distribution='logitnormal', device=device)
xt = noise * sigma_t + clean_latent * (1 - sigma_t)
xt = clean_latent * cond_mask + xt * (1 - cond_mask)
in_timestep = cond_indicator * 0.0001 + (1 - cond_indicator) * sigma_t
pred_velocity = dit(
    hidden_states=xt,
    condition_mask=cond_mask,
    timestep=in_timestep,
    encoder_hidden_states=prompt_embeds,
    padding_mask=padding_mask,
    return_dict=False,
    target_velocity=noise - clean_latent
)
pred_velocity = target_velocity * cond_mask + pred_velocity * (1 - cond_mask)
loss = F.mse_loss(pred_velocity.float(), target_velocity.float())

Optimiseur et Planificateur

Nous utilisons torch.optim.AdamW comme optimiseur et get_linear_schedule_with_warmup de diffusers.optimization comme planificateur. Le planificateur augmente linéairement le taux d'apprentissage sur scheduler_warm_up_steps, atteint un pic à scheduler_f_max × learning_rate, puis décroit linéairement à scheduler_f_min × learning_rate sur les étapes d'entraînement restantes.

lora_params = [p for p in dit.parameters() if p.requires_grad]
optimizer = torch.optim.AdamW(lora_params, lr=args.learning_rate, weight_decay=args.weight_decay)

lr_scheduler = get_linear_schedule_with_warmup(
    num_warmup_steps=args.scheduler_warm_up_steps,
    num_training_steps=args.num_training_steps,
    f_min=args.scheduler_f_min,
    f_max=args.scheduler_f_max,
)

Les poids LoRA sont sauvegardés dans le format diffusers tous les args.checkpointing_epochs époques :

if (epoch+1) % args.checkpointing_epochs == 0:
    if accelerator.is_main_process:
        save_path = os.path.join(args.output_dir, f"checkpoint-{epoch}")
        accelerator.save_state(save_path)

accelerator.save_state() écrit un fichier pytorch_lora_weights.safetensors dans save_path, qui est le fichier d'adaptateur que vous passerez au pipeline lors de l'inférence.

Commande d'entraînement

Utilisez le script shell fourni comme point de départ :

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

accelerate launch --mixed_precision="bf16" train_cosmos_predict25_lora.py \
--pretrained_model_name_or_path=$MODEL_NAME \
--revision diffusers/base/post-trained \
--train_data_dir=$DATA_DIR \
--train_batch_size=1 \
--num_train_epochs=500 \
--checkpointing_epochs=100 \
--output_dir=$OUT_DIR \
--report_to=wandb \
--height 432 --width 768 \
--allow_tf32 --gradient_checkpointing \
--lora_rank $lora_rank --lora_alpha $lora_rank

lora_rank contrôle le rang de la décomposition de faible rang. Un rang plus élevé signifie plus de paramètres entraînables et une plus grande capacité d'expression, au prix d'une consommation de mémoire plus importante et d'un fichier d'adaptateur plus volumineux. Nous utilisons rank=32 comme point de départ, ce qui donne environ 50 millions de paramètres entraînables.

lora_alpha est un facteur d'échelle appliqué à la mise à jour LoRA : la variation de poids est mise à l'échelle par lora_alpha / lora_rank avant d'être ajoutée aux poids de base gelés. En définissant lora_alpha = lora_rank (comme ici), ce facteur d'échelle reste à 1.0, de sorte que la mise à jour LoRA est appliquée à pleine puissance sans aucun amortissement supplémentaire.

Pour utiliser DoRA au lieu de LoRA, ajoutez --use_dora à la commande.

Pour l'entraînement multi-GPU, accelerate gère automatiquement la distribution. Empiriquement, nous constatons qu'un entraînement de 100 époques donne déjà des résultats décents sur cette tâche, ce qui prend 17 heures sur un seul H100 et 2,5 heures sur 8 GPU H100.

Exécution de l'inférence avec votre LoRA

Une fois l'entraînement terminé, utilisez eval_cosmos_predict25_lora.py pour générer des vidéos à partir de l'ensemble de données d'évaluation. Le script lit les fichiers .png et .txt appariés depuis gr1_dataset/test, génère une vidéo pour chacun, et écrit des fichiers .mp4 dans --output_dir.

ImageDataset lit le fichier .txt en une chaîne d'invite et utilise load_image de diffusers.utils pour charger le .png en tant que PIL.Image.Image :

def __getitem__(self, idx):
    img_path, txt_path, stem = self.samples[idx]
    image = load_image(img_path)
    with open(txt_path) as f:
        [prompt](/glossaire/prompt) = f.read().strip()
    return {"image": image, "prompt": prompt, "stem": stem}

Chargement du Pipeline et des poids LoRA/DoRA

from diffusers import Cosmos2_5_PredictBasePipeline

pipe = Cosmos2_5_PredictBasePipeline.from_pretrained(
    "nvidia/Cosmos-Predict2.5-2B",
    revision="diffusers/base/post-trained",
    device_map="cuda",
    torch_dtype=torch.bfloat16,
)

pipe.load_lora_weights("/path/to/lora/checkpoint")
pipe.fuse_lora(lora_scale=1.0)

fuse_lora fusionne les poids de l'adaptateur dans le modèle de base, éliminant tout surcoût d'inférence dû à la décomposition LoRA/DoRA.

Génération de bruit latent initial

Pour garantir la reproductibilité, la fonction arch_invariant_rand génère le bruit latent initial via NumPy, rendant le bruit invariant aux architectures GPU. Si la reproductibilité n'est pas une préoccupation, les utilisateurs n'ont pas besoin de fournir de bruit d'entrée au pipeline.

latent_shape = pipe.get_latent_shape_cthw(args.height, args.width, args.num_output_frames)