Fine‑Tuning NVIDIA Cosmos Predict 2.5 with LoRA/DoRA for Robot Video Generation
Use LoRA/DoRA adapters to fine‑tune Cosmos Predict 2.5 on a single GPU, reducing memory and keeping adapters portable.
Bump your training pipeline to use LoRA/DoRA adapters, set lora_rank=32, and run accelerate launch with bf16; monitor training with wandb.
Summary
NVIDIA Cosmos Predict 2.5, a 2 B‑parameter world model, can generate physically plausible videos conditioned on text, images, or video clips. To adapt it to domain‑specific tasks such as robot manipulation, the guide shows how to fine‑tune the model using LoRA or DoRA adapters, which inject small trainable modules into the frozen base and keep adapter files portable. The training pipeline relies on diffusers, accelerate, and peft, and requires at least an 80 GB GPU for single‑GPU runs, with 8 × H100s recommended for faster iteration. LoRA adapters are added to the DiT attention projections (to_q, to_k, to_v, to_out.0) and feed‑forward layers (ff.net.0.proj, ff.net.2), and the model is trained with a rectified‑flow loss that predicts velocity toward clean data. The optimizer is AdamW with a linear warm‑up scheduler, and checkpoints are saved as safetensors every few epochs. Inference uses the eval_cosmos_predict25_lora.py script to generate synthetic robot trajectories from paired prompt‑image datasets. The guide also details how to enable DoRA by passing --use_dora, which decomposes weights into magnitude and direction before applying the low‑rank update. Overall, the approach allows teams to fine‑tune Cosmos Predict 2.5 on a single GPU while preserving the model’s general knowledge and producing domain‑specific synthetic videos for downstream robot learning.
Key changes
- LoRA/DoRA inject small trainable adapters into frozen base model
- Supports single‑ and multi‑GPU training via diffusers and accelerate
- Requires at least 80 GB GPU for single‑GPU runs; 8× H100 recommended
- Adapters added to DiT attention projections and feed‑forward layers
- Training uses rectified‑flow loss predicting velocity toward clean data
- Optimizer AdamW with linear warm‑up scheduler; checkpoints saved as safetensors
- Inference uses eval_cosmos_predict25_lora.py to generate synthetic videos