动机
NVIDIA Cosmos Predict 2.5 是一个大规模世界模型，能够根据文本、图像或视频片段生成物理上合理的视频。为了将其适配到特定领域，例如机器人操作或特定摄像机视角，团队仍然需要进行针对性的微调。
训练机器人策略需要演示数据，但收集真实机器人轨迹既缓慢又昂贵。使用微调后的视频世界模型生成合成轨迹提供了一种可扩展的替代方案。然而，对 20 亿参数的模型进行全量微调成本高昂，并且存在灾难性遗忘通用知识的风险。LoRA 和 DoRA 将小型可训练适配器模块注入冻结的基础模型中，降低了内存需求，同时保持了适配器文件的小巧和便携性。这使得在单张 GPU 上进行微调变得可行，并且可以在推理时灵活地为不同领域切换适配器。
本指南将逐步介绍如何使用 LoRA 和 DoRA 对 Cosmos Predict 2.5 进行参数高效微调，利用 diffusers 和 accelerate 库，并支持单 GPU 和多 GPU 训练。然后，我们将展示如何使用微调后的模型为下游机器人学习任务生成合成机器人轨迹。
环境要求
- Python 3.10+
- 配备 CUDA 的 PyTorch 2.5+
- diffusers（会自动引入 transformers 和 peft）、accelerate
- 可选：安装 wandb 以监控训练过程
- 至少需要一张 80 GB 的 GPU 用于单 GPU 训练；建议使用 8× H100 以加快迭代速度
在你的机器上安装依赖项：
pip install -U "diffusers[torch]" transformers accelerate peft wandb
准备数据
安装 diffusers 后，导航至 examples/cosmos 目录以查看示例代码。
我们使用与 GR00T Dreams 后训练方案相同的数据集：
- 训练数据集：92 个机器人操作视频，附带描述拾取和放置任务的文本提示词。
- 测试数据集：50 个（提示词，图像）对。模型应根据输入的文本提示词和初始帧图像生成一个视频。
使用 download_and_preprocess_datasets.sh 脚本下载并预处理训练和测试数据集：
bash download_and_preprocess_datasets.sh
生成的训练数据集文件夹结构如下：
gr1_dataset/train
├── metas/
│ └── *.txt
├── videos/
│ └── *.mp4
└── metadata.csv
评估数据集是一个扁平目录，包含配对的 .txt 和 .png 文件，用于（提示词，图像）对：
gr1_dataset/test
├── filename1.txt
├── filename1.png
├── filename2.txt
├── filename2.png
└── ...
训练
在本节中，我们将逐步讲解 `train_cosmos_predict25_lora.py` 中的实现过程。
视频数据集
`VideoDataset` 从 `args.train_data_dir`（在我们的示例中为 `gr1_dataset/train`）加载每个样本，作为一对（描述文本，视频）。对于长度超过 `args.num_frames` 的视频，它会在每个训练周期随机采样一个连续的 `args.num_frames` 帧窗口，从而实现时间维度的数据增强。在内部，`diffusers.video_processor` 中的 `VideoProcessor` 会将原始帧调整大小并归一化为形状为（通道数，帧数，高度，宽度）的张量。
train_dataset = VideoDataset(
dataset_dir=args.train_data_dir,
num_frames=args.num_frames,
video_size=[args.height, args.width],
)
初始化适配器
Cosmos Predict 2.5 由三个子模块组成：
- 一个将视频编码为潜在向量的 VAE
- 一个将文本提示词编码为提示词嵌入向量的文本编码器
- 用于在潜在空间中进行扩散的 DiT
在训练期间，所有 VAE、文本编码器和 DiT 的权重均被冻结。LoRA 适配器被注入到 DiT 的注意力投影层（`to_q`、`to_k`、`to_v`、`to_out.0`）和前馈网络层（`ff.net.0.proj`、`ff.net.2`）中。然后，可训练的 LoRA 参数会被向上转换为 float32 类型，以在 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,
)
# freeze all base weights
dit = pipe.transformer
vae = pipe.vae
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, # set True to switch to DoRA
)
dit.add_adapter(lora_config)
cast_training_params(dit, dtype=torch.float32) # LoRA params in fp32
传入 `use_dora=True` 会切换为 DoRA，该方法在应用低秩更新之前，将每个权重分解为幅度和方向。训练循环的其他部分无需任何更改。
损失函数
Cosmos Predict 2.5 使用修正流：模型被训练来预测一个速度，该速度能够线性地将噪声样本向原始“干净”数据迁移。具体来说，在时间步 t，会在一个采样的噪声水平 σt 下构建一个带噪声的插值 `xt = σt·噪声 + (1−σt)·干净数据`，模型则通过均方误差损失来学习预测目标速度 `噪声 − 干净数据`。视频的前两帧被用作条件，因此不会向它们的潜在向量添加噪声。
训练损失遵循 Cosmos Predict 2.5 所使用的修正流公式：
# Sample timestep with logit-normal distribution
sigma_t = sample_train_sigma_t(bsz, distribution='logitnormal', device=device)
# Rectified flow interpolates between clean latent and noise
xt = noise * sigma_t + clean_latent * (1 - sigma_t)
# Conditional generation: DiT conditions on the first two frames of the video, the timestep, and the prompt embeds
# `cond_indicator` and `cond_mask` have values = 1 for the first two frames and 0 for other frames
xt = clean_latent * cond_mask + xt * (1 - cond_mask)
in_timestep = cond_indicator * 0.0001 + (1 - cond_indicator) * sigma_t
# Forward
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,
)[0]
# MSE loss is computed only on the non-conditioned frames
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())
优化器和调度器
我们使用 `torch.optim.AdamW` 作为优化器，并使用 `diffusers.optimization` 中的 `get_linear_schedule_with_warmup` 作为调度器。该调度器会在 `scheduler_warm_up_steps` 步内线性预热学习率，使其达到峰值 `scheduler_f_max × learning_rate`，然后在剩余的 `num_training_steps` 步中线性衰减至 `scheduler_f_min × learning_rate`。
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(
optimizer,
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,
)
检查点保存
LoRA 权重会按照 `args.checkpointing_epochs` 参数设定的轮次间隔，以 diffusers 格式保存。
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()` 会在 `save_path` 路径下写入一个 `pytorch_lora_weights.safetensors` 文件，该文件即为推理时需传入 pipeline 的适配器文件。
训练命令
以下面的 shell 脚本作为起点：
export MODEL_NAME="nvidia/Cosmos-Predict2.5-2B"
export DATA_DIR="gr1_dataset/train"
export OUT_DIR=YOUR_OUTPUT_DIR
lora_rank=32
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 \
--seed=0 \
--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` 控制低秩分解的秩。秩越高，可训练参数越多，表达能力越强，但代价是内存占用更大、适配器文件也更大。我们以 rank=32 为起点，可训练参数约为 5000 万。
`lora_alpha` 是应用于 LoRA 更新的缩放因子：权重增量在加到冻结的基础权重上之前，会先乘以 `lora_alpha / lora_rank`。将 `lora_alpha` 设为与 `lora_rank` 相等（如本例所示），可使该缩放因子保持为 1.0，从而让 LoRA 更新以完整强度生效，无需任何额外衰减。
若要使用 DoRA 而非 LoRA，请在命令中添加 `--use_dora`。
对于多 GPU 训练，accelerate 会自动处理分布式逻辑。根据经验，我们发现在该任务上训练 100 个 epoch 即可获得不错的结果，在单张 H100 上耗时 17 小时，在 8 张 H100 GPU 上耗时 2.5 小时。
使用你的 LoRA 运行推理
训练完成后，使用 `eval_cosmos_predict25_lora.py` 脚本从评估数据集中生成视频。该脚本会读取 `gr1_dataset/test` 目录下配对的 `.png` 和 `.txt` 文件，为每个样本生成一段视频，并将 `.mp4` 文件写入 `--output_dir` 指定的目录。
ImageDataset
ImageDataset 会将 `.txt` 文件读取为提示词字符串，并使用 `diffusers.utils` 中的 `load_image` 函数将 `.png` 文件加载为 `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 = f.read().strip()
return {"image": image, "prompt": prompt, "stem": stem}
加载 Pipeline 与 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` 会将适配器权重合并到基础模型中，从而消除 LoRA/DoRA 分解带来的任何推理开销。
生成初始潜变量噪声
为确保可复现性，`arch_invariant_rand` 函数通过 NumPy 生成初始潜变量噪声，使噪声不受 GPU 架构影响。如果不需要可复现性，用户无需向 pipeline 提供输入噪声。
# generation starts from random noise with the same shape as the latent
latent_shape = pipe.get_latent_shape_cthw(args.height, args.width, args.num_output_frames)
noises = arch_invariant_rand(
(args.batch_size, *latent_shape), dtype=torch.float32, device=args.device, seed=args.seed
)
frames = pipe(
image=image, # PIL Image: the conditioning first frame
prompt=prompt,
num_frames=args.num_output_frames,
num_inference_steps=args.num_steps,
height=args.height,
width=args.width,
latents=noises, # optional
).frames[0]
export_to_video(frames, "output.mp4", fps=16)
推理命令
export LORA_DIR=YOUR_ADAPTER_DIR
export DATA_DIR="gr1_dataset/test"
export OUT_DIR=YOUR_EVAL_OUTPUT_DIR
python eval_cosmos_predict25_lora.py \
--data_dir $DATA_DIR \
--output_dir $OUT_DIR \
--lora_dir $LORA_DIR \
--height 432 --width 768 \
--num_output_frames 93 \
--num_steps 36 \
--seed 0
若要评估不包含任何 LoRA 的基础模型，请省略 `--lora_dir` 参数。
评估指标
Sampson 误差
Sampson 误差是一种几何误差度量指标，用于衡量匹配关键点到其对应极线之间的距离。在生成视频的语境下，较低的 Sampson 误差意味着帧与帧之间（或不同相机视角之间）的运动在几何上是一致的。数值越高，则表明存在抖动、模型幻觉产生的运动或多视角不一致的问题。
我们遵循 Cosmos Predict 评估指南，使用以下两个指标来评估生成视频的几何质量：
- 时间维度 Sampson 误差：在单个相机视角内，对连续帧之间进行计算，用于衡量时间稳定性。
- 跨视角 Sampson 误差：在不同相机视角的同步帧之间进行计算，用于衡量多视角几何对齐程度。
大语言模型作为评判者
我们使用 Cosmos Reason2 作为大语言模型评判者，对每个样本进行 1 到 5 分的评分。我们设计了两个评分标准：
- 物理合理性（video_physics.yaml）：评判者在未看到文本提示词的情况下，评估视频是否符合物理常识。
- 指令遵循（video_IF.yaml）：评判者将提示词和视频同时作为输入，评估所描述的任务是否被正确完成。
video_physics.yaml
system_prompt: "You are a helpful assistant."
user_prompt: |
You are a helpful video analyzer. Evaluate whether the video follows physical commonsense.
Evaluation Criteria:
1. **Object Behavior:** Do objects behave according to their expected physical properties (e.g., rigid objects do not deform unnaturally, fluids flow naturally)?
2. **Motion and Forces:** Are motions and forces depicted in the video consistent with real-world physics (e.g., gravity, inertia, conservation of momentum)?
3. **Interactions:** Do objects interact with each other and their environment in a plausible manner (e.g., no unnatural penetration, appropriate reactions on impact)?
4. **Consistency Over Time:** Does the video maintain consistency across frames without abrupt, unexplainable changes in object behavior or motion?
Instructions for Scoring:
- **1:** No adherence to physical commonsense. The video contains numerous violations of fundamental physical laws.
- **2:** Poor adherence. Some elements follow physics, but major violations are present.
- **3:** Moderate adherence. The video follows physics for the most part but contains noticeable inconsistencies.
- **4:** Good adherence. Most elements in the video follow physical laws, with only minor issues.
- **5:** Perfect adherence. The video demonstrates a strong understanding of physical commonsense with no violations.
Does this video adhere to the physical laws?
video_IF.yaml
system_prompt: "You are a helpful assistant."
user_prompt: |
You are a helpful video analyzer. Evaluate whether the video follows the given instruction.
Instruction: {instruction}
Evaluation Criteria:
1. **Task Completion:** Does the video show the task described in the instruction being completed?
2. **Action Accuracy:** Are the actions performed in the video consistent with what the instruction specifies?
3. **Object Interaction:** Does the robot or agent interact with the correct objects as described in the instruction?
4. **Goal Achievement:** Is the final state of the video consistent with the expected outcome of the instruction?
5. **Correct Hand Usage:** Does the video show the correct hand performing the action?
Instructions for Scoring:
- **1:** No adherence to the instruction. The video shows actions completely unrelated to the instruction.
- **2:** Poor adherence. Some elements match the instruction, but major deviations are present.
- **3:** Moderate adherence. The video follows the instruction for the most part but contains noticeable deviations.
- **4:** Good adherence. Most elements in the video match the instruction, with only minor issues.
- **5:** Perfect adherence. The video fully follows the instruction with no deviations.
Does this video follow the instruction?
结果
定性分析
我们比较了基础模型（微调前）、LoRA 和 DoRA 在测试集前两个样本上生成的视频。
提示词：用左手将深绿色黄瓜从圆形灰色垫子上拿起，放到浅绿色碗的上方。
| 训练前 | LoRA r=32 | DoRA r=32 |
|---|---|---|
提示词：用右手将橙汁盒从粉色盘子中央拿起，放到绿色碗的中央。
| 训练前 | LoRA r=32 | DoRA r=32 |
|---|---|---|
在微调之前，基础模型在多个方面表现不佳：机器手属于分布外数据，导致模型在后续帧中幻觉生成了人类的手；模型无法可靠地使用提示词中指定的正确手部；并且生成的视频存在明显的抖动。使用 LoRA 和 DoRA 进行微调解决了上述所有三个问题。
定量分析
我们在不同设置下微调了四个适配器：秩为 8 和 32 的 LoRA 与 DoRA。对于每个测试样本，我们使用不同种子生成 5 个视频，并报告所有种子上的平均得分，采用评估指标部分介绍的三种指标。
萨姆森误差（越低越好）。微调后，时序和跨视图萨姆森误差均有所下降，表明时序稳定性和多视图几何一致性得到改善。
物理合理性得分（越高越好）。与基础模型相比，微调后的模型生成的视频更符合物理常识。
指令遵循得分（越高越好）。微调后的模型能更可靠地完成提示词中描述的任务，包括使用正确的手部以及与指定物体进行交互。
结论：训练 100 个 epoch（在 8× H100 上约需 2.5 小时）已足以显著提升所有三项指标。LoRA 和 DoRA 均收敛到相似性能，证实 DoRA 中额外的幅度-方向分解不会造成损害，且在极低秩下可能有所帮助，但在此场景下并非必需。
更大的秩（32 对比 8）能提升指令遵循能力（模型有更多容量来精确学习使用哪只手以及交互哪些物体），但不会改善几何一致性或物理合理性。我们推测这是因为几何和物理先验知识主要由世界模型的冻结权重捕获；LoRA 适配器仅需将分布向域内机器人外观和任务结构偏移，这在秩为 8 时即可实现。
何时使用 DoRA 对比 LoRA：如果内存非常紧张或适配器文件大小很重要，请从 LoRA r=8 开始。如果预算充足且观察到 LoRA 在低秩下训练不稳定，DoRA r=32 是一个合理的替代方案，因为幅度-方向分解有助于稳定学习。
- 请访问我们的 Cosmos Cookbook，获取构建、适配和部署 Cosmos WFM 的分步工作流程、技术方案和具体示例。
- 在 Hugging Face 和 GitHub 上探索新的开源 Cosmos 模型和数据集，或在 build.nvidia.com 上试用模型。
- 加入社区，进入我们的 Cosmos Discord 频道。
- 已经在使用 Cosmos 了吗？了解更多关于如何贡献的信息。
Motivation
NVIDIA Cosmos Predict 2.5 is a large-scale world model capable of generating physically plausible videos conditioned on text, images, or video clips. To adapt it to a specific domain, such as robot manipulation or a particular camera viewpoint, teams still need targeted fine-tuning.
Training robot policies requires demonstration data, but collecting real-robot trajectories is slow and expensive. Generating synthetic trajectories with a fine-tuned video world model offers a scalable alternative. However, full fine-tuning of a 2B-parameter model is expensive and risks catastrophic forgetting of general knowledge. LoRA and DoRA inject small trainable adapter modules into the frozen base model, reducing memory requirements while keeping the adapter files small and portable. This makes it practical to fine-tune on a single GPU and flexibly swap adapters for different domains at inference.
This guide walks through parameter-efficient fine-tuning of Cosmos Predict 2.5 with LoRA and DoRA, using the diffusers and accelerate libraries with support for both single- and multi-GPU training. We then show how to use the fine-tuned model to generate synthetic robot trajectories for downstream robot learning tasks.
Requirements
- Python 3.10+
- PyTorch 2.5+ with CUDA
diffusers(pulls intransformersandpeftautomatically),accelerate- Optional: install
wandbto monitor training - At minimum one 80 GB GPU for single-GPU training; 8× H100s recommended for faster iteration
Install dependencies on your machine:
pip install -U "diffusers[torch]" transformers accelerate peft wandb
Preparing Data
After installing diffusers, navigate to examples/cosmos to explore the example code.
We use the same datasets as the GR00T Dreams post-training recipe:
- Training Dataset: 92 robot manipulation videos with text prompts describing pick-and-place tasks.
- Test Dataset: 50 (prompt, image) pairs. The model should generate a video based on the input text prompt and the initial frame image.
Download and preprocess the training and test datasets using download_and_preprocess_datasets.sh:
bash download_and_preprocess_datasets.sh
The resulting training dataset folder looks like this:
gr1_dataset/train
├── metas/
│ └── *.txt
├── videos/
│ └── *.mp4
└── metadata.csv
The eval dataset is a flat directory of paired .txt and .png files for the (prompt, image) pairs:
gr1_dataset/test
├── filename1.txt
├── filename1.png
├── filename2.txt
├── filename2.png
└── ...
Training
In this section, we walk through the implementation in train_cosmos_predict25_lora.py.
VideoDataset
VideoDataset loads each sample as a (caption, video) pair from args.train_data_dir (gr1_dataset/train in our example). For videos longer than args.num_frames, it samples a random contiguous window of args.num_frames each epoch, enabling temporal augmentation. Internally, VideoProcessor from diffusers.video_processor resizes and normalizes the raw frames into a tensor of shape (channels, frames, height, width).
train_dataset = VideoDataset(
dataset_dir=args.train_data_dir,
num_frames=args.num_frames,
video_size=[args.height, args.width],
)
Initialize Adapter
Cosmos Predict 2.5 consists of three submodules:
- A VAE that encodes videos into latents
- A text encoder that encodes text prompts into prompt embeddings
- DiT for diffusion in the latent space
During training, all VAE, text encoder, and DiT weights are frozen. LoRA adapters are injected into the DiT's attention projections (to_q, to_k, to_v, to_out.0) and feedforward layers (ff.net.0.proj, ff.net.2). The trainable LoRA parameters are then upcast to float32 for numerical stability under bf16 mixed precision.
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,
)
# freeze all base weights
dit = pipe.transformer
vae = pipe.vae
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, # set True to switch to DoRA
)
dit.add_adapter(lora_config)
cast_training_params(dit, dtype=torch.float32) # LoRA params in fp32
Passing use_dora=True switches to DoRA, which decomposes each weight into magnitude and direction before applying the low-rank update. No other changes to the training loop are needed.
Loss
Cosmos Predict 2.5 uses rectified flow: the model is trained to predict the velocity that linearly transports a noise sample toward the original "clean" data. Concretely, at timestep t, a noisy interpolation xt = σt·noise + (1−σt)·clean is constructed at a sampled noise level σt, and the model learns to predict the target velocity noise − clean via the mean-squared errors (MSE loss). The first two frames of the video are used as conditioning, and thus no noise is added to their latents..
The training loss follows the rectified flow formulation used by Cosmos Predict 2.5:
# Sample timestep with logit-normal distribution
sigma_t = sample_train_sigma_t(bsz, distribution='logitnormal', device=device)
# Rectified flow interpolates between clean latent and noise
xt = noise * sigma_t + clean_latent * (1 - sigma_t)
# Conditional generation: DiT conditions on the first two frames of the video, the timestep, and the prompt embeds
# `cond_indicator` and `cond_mask` have values = 1 for the first two frames and 0 for other frames
xt = clean_latent * cond_mask + xt * (1 - cond_mask)
in_timestep = cond_indicator * 0.0001 + (1 - cond_indicator) * sigma_t
# Forward
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,
)[0]
# MSE loss is computed only on the non-conditioned frames
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())
Optimizer and Scheduler
We use torch.optim.AdamW as the optimizer and get_linear_schedule_with_warmup from diffusers.optimization as the scheduler. The scheduler linearly warms up the learning rate over scheduler_warm_up_steps, peaks at scheduler_f_max × learning_rate, then linearly decays to scheduler_f_min × learning_rate over the remaining num_training_steps.
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(
optimizer,
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,
)
Checkpointing
LoRA weights are saved in the diffusers format every args.checkpointing_epochs epochs:
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() writes a pytorch_lora_weights.safetensors file to save_path, which is the adapter file you will pass to the pipeline at inference time.
Training Command
Use the provided shell script as a starting point:
export MODEL_NAME="nvidia/Cosmos-Predict2.5-2B"
export DATA_DIR="gr1_dataset/train"
export OUT_DIR=YOUR_OUTPUT_DIR
lora_rank=32
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 \
--seed=0 \
--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 controls the rank of the low-rank decomposition. A higher rank means more trainable parameters and greater expressive capacity, at the cost of more memory and a larger adapter file. We use rank=32 as a starting point, resulting in ~50M trainable parameters.
lora_alpha is a scaling factor applied to the LoRA update: the weight delta is scaled by lora_alpha / lora_rank before being added to the frozen base weights. Setting lora_alpha = lora_rank (as done here) keeps this scale factor at 1.0, so the LoRA update is applied at full strength without any additional dampening.
To use DoRA instead of LoRA, add --use_dora to the command.
For multi-GPU training, accelerate handles the distribution automatically. Empirically, we find that training with 100 epochs already yields decent results on this task, which takes 17 hours on a single H100 and 2.5 hours on 8 H100 GPUs.
Running Inference with Your LoRA
Once training is complete, use eval_cosmos_predict25_lora.py to generate videos from the eval dataset. The script reads paired .png and .txt files from gr1_dataset/test, generates a video for each, and writes .mp4 files to --output_dir.
ImageDataset
ImageDataset reads the .txt file into a prompt string and uses load_image from diffusers.utils to load the .png as a 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 = f.read().strip()
return {"image": image, "prompt": prompt, "stem": stem}
Loading the Pipeline and LoRA/DoRA Weights
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 merges the adapter weights into the base model, eliminating any inference overhead from the LoRA/DoRA decomposition.
Generating initial latent noise
To ensure reproducibility, the arch_invariant_rand function generates the initial latent noise via NumPy, making the noise invariant to GPU architectures. If reproducibility is not a concern, users do not need to provide input noise to the pipeline.
# generation starts from random noise with the same shape as the latent
latent_shape = pipe.get_latent_shape_cthw(args.height, args.width, args.num_output_frames)
noises = arch_invariant_rand(
(args.batch_size, *latent_shape), dtype=torch.float32, device=args.device, seed=args.seed
)
frames = pipe(
image=image, # PIL Image: the conditioning first frame
prompt=prompt,
num_frames=args.num_output_frames,
num_inference_steps=args.num_steps,
height=args.height,
width=args.width,
latents=noises, # optional
).frames[0]
export_to_video(frames, "output.mp4", fps=16)
Inference Command
export LORA_DIR=YOUR_ADAPTER_DIR
export DATA_DIR="gr1_dataset/test"
export OUT_DIR=YOUR_EVAL_OUTPUT_DIR
python eval_cosmos_predict25_lora.py \
--data_dir $DATA_DIR \
--output_dir $OUT_DIR \
--lora_dir $LORA_DIR \
--height 432 --width 768 \
--num_output_frames 93 \
--num_steps 36 \
--seed 0
To evaluate the base model without any LoRA, omit --lora_dir.
Evaluation Metrics
Sampson Error
Sampson Error is a geometric error metric that measures the distance from matched keypoints to their corresponding epipolar lines. In the context of generated video, a low Sampson error means the motion between frames (or between camera views) is geometrically consistent. Higher values indicate jitter, hallucinated motion, or multi-view inconsistencies.
We follow the Cosmos Predict evaluation guide and evaluate the geometric quality of generated videos using two metrics:
- Temporal Sampson Error: computed between consecutive frames within a single camera view, measuring temporal stability.
- Cross-view Sampson Error: computed between simultaneous frames from different camera views, measuring multi-view geometric alignment.
LLM-as-a-Judge
We use Cosmos Reason2 as an LLM judge, scoring each example from 1 to 5. We design two rubrics:
- Physical plausibility (video_physics.yaml): the judge evaluates whether the video obeys physical commonsense, without seeing the text prompt.
- Instruction following (video_IF.yaml): the judge takes both the prompt and the video as input and evaluates whether the described task is completed correctly.
video_physics.yaml
system_prompt: "You are a helpful assistant."
user_prompt: |
You are a helpful video analyzer. Evaluate whether the video follows physical commonsense.
Evaluation Criteria:
1. **Object Behavior:** Do objects behave according to their expected physical properties (e.g., rigid objects do not deform unnaturally, fluids flow naturally)?
2. **Motion and Forces:** Are motions and forces depicted in the video consistent with real-world physics (e.g., gravity, inertia, conservation of momentum)?
3. **Interactions:** Do objects interact with each other and their environment in a plausible manner (e.g., no unnatural penetration, appropriate reactions on impact)?
4. **Consistency Over Time:** Does the video maintain consistency across frames without abrupt, unexplainable changes in object behavior or motion?
Instructions for Scoring:
- **1:** No adherence to physical commonsense. The video contains numerous violations of fundamental physical laws.
- **2:** Poor adherence. Some elements follow physics, but major violations are present.
- **3:** Moderate adherence. The video follows physics for the most part but contains noticeable inconsistencies.
- **4:** Good adherence. Most elements in the video follow physical laws, with only minor issues.
- **5:** Perfect adherence. The video demonstrates a strong understanding of physical commonsense with no violations.
Does this video adhere to the physical laws?
video_IF.yaml
system_prompt: "You are a helpful assistant."
user_prompt: |
You are a helpful video analyzer. Evaluate whether the video follows the given instruction.
Instruction: {instruction}
Evaluation Criteria:
1. **Task Completion:** Does the video show the task described in the instruction being completed?
2. **Action Accuracy:** Are the actions performed in the video consistent with what the instruction specifies?
3. **Object Interaction:** Does the robot or agent interact with the correct objects as described in the instruction?
4. **Goal Achievement:** Is the final state of the video consistent with the expected outcome of the instruction?
5. **Correct Hand Usage:** Does the video show the correct hand performing the action?
Instructions for Scoring:
- **1:** No adherence to the instruction. The video shows actions completely unrelated to the instruction.
- **2:** Poor adherence. Some elements match the instruction, but major deviations are present.
- **3:** Moderate adherence. The video follows the instruction for the most part but contains noticeable deviations.
- **4:** Good adherence. Most elements in the video match the instruction, with only minor issues.
- **5:** Perfect adherence. The video fully follows the instruction with no deviations.
Does this video follow the instruction?
Results
Qualitative Analysis
We compare videos generated by the base model (before fine-tuning), LoRA, and DoRA on the first two examples from the test set.
Prompt: Use the left hand to pick up dark green cucumber from on circular gray mat to above beige bowl.
| Before Training | LoRA r=32 | DoRA r=32 |
|---|---|---|
Prompt: Use the right hand to pick up orange juice carton from center of pink plate to center of green bowl.
| Before Training | LoRA r=32 | DoRA r=32 |
|---|---|---|
Before fine-tuning, the base model struggles in several ways: robot hands are out-of-distribution, causing the model to hallucinate human hands in later frames; it does not reliably use the correct hand specified in the prompt; and the generated videos exhibit noticeable jitter. Fine-tuning with LoRA and DoRA addresses all three issues.
Quantitative Analysis
We fine-tune four adapters under different settings: LoRA and DoRA with rank 8 and 32. For each test example, we generate 5 videos with different seeds and report the average score across seeds, using the three metrics introduced in the Evaluation Metrics section.
Sampson Error (lower is better). Both Temporal and Cross-view Sampson Errors decrease after fine-tuning, indicating improved temporal stability and multi-view geometric consistency.
Physical plausibility score (higher is better). Fine-tuned models generate videos that better adhere to physical commonsense compared to the base model.
Instruction following score (higher is better). Fine-tuned models more reliably complete the task described in the prompt, including using the correct hand and interacting with the specified objects.
Conclusion: Training for 100 epochs (~2.5 hours on 8× H100s) is already sufficient to substantially improve all three metrics. Both LoRA and DoRA converge to similar performance, confirming that the extra magnitude-direction decomposition in DoRA does not hurt and may help at very low ranks, but is not necessary here.
Larger rank (32 vs 8) boosts instruction following (the model has more capacity to learn precisely which hand to use and which objects to interact with), but does not improve geometric consistency or physical plausibility. We hypothesize that this is because geometric and physical priors are largely captured by the world model's frozen weights; the LoRA adapter only needs to shift the distribution toward in-domain robot appearance and task structure, which is achievable at rank 8.
When to use DoRA vs LoRA: If memory is very tight or adapter file size matters, start with LoRA r=8. If you have budgets and observe training instability with LoRA at low rank, DoRA r=32 is a reasonable alternative, as the magnitude–direction decomposition can help stabilize learning.
- Visit our Cosmos Cookbook for step-by-step workflows, technical recipes, and concrete examples for building, adapting, and deploying Cosmos WFMs.
- Explore new open Cosmos models and datasets on Hugging Face and GitHub or try models on build.nvidia.com.
- Be part of the community and join our Cosmos Discord channel.
- Already using Cosmos? Learn more about how to contribute.