使用 LoRA/DoRA 微调 NVIDIA Cosmos Predict 2.5 以生成机器人视频

来源：https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation 返回文章列表 (https://huggingface.co/blog)

- 动机 (https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#motivation) - 需求 (https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#requirements) - 数据准备 (https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#preparing-data) - 训练 (https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#training)- VideoDataset (https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#videodataset) - 初始化适配器 (https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#initialize-adapter) - 损失函数 (https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#loss) - 优化器与调度器 (https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#optimizer-and-scheduler) - 检查点 (https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#checkpointing) - 训练命令 (https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#training-command) - 使用你的 LoRA 进行推理 (https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#running-inference-with-your-lora)- ImageDataset (https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#imagedataset) - 加载流水线和 LoRA/DoRA 权重 (https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#loading-the-pipeline-and-loradora-weights) - 生成初始潜变量噪声 (https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#generating-initial-latent-noise) - 推理命令 (https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#inference-command) - 评估指标 (https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#evaluation-metrics)- Sampson 误差 (https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#sampson-error) - LLM 作为裁判 (https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#llm-as-a-judge) - 结果 (https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#results)- 定性分析 (https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#qualitative-analysis) - 定量分析 (https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#quantitative-analysis) https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#motivation动机

NVIDIA Cosmos Predict 2.5 (https://arxiv.org/abs/2511.00062) 是一个大规模世界模型 (https://www.nvidia.com/en-us/glossary/world-models/?ncid=so-nvsh-876275)，能够根据文本、图像或视频片段生成物理上合理的视频。为了将其适应到特定领域，例如机器人操作或特定的摄像头视角，团队仍然需要进行针对性的微调。

训练机器人策略需要演示数据，但收集真实的机器人轨迹既慢又昂贵。使用微调后的视频世界模型生成合成轨迹提供了一种可扩展的替代方案。然而，对 2B 参数模型进行全参数微调成本高昂，并且存在灾难性遗忘通用知识的风险。LoRA (https://arxiv.org/abs/2106.09685) 和 DoRA (https://arxiv.org/abs/2402.09353) 将小型可训练适配器模块注入冻结的基础模型中，降低了内存需求，同时保持适配器文件小巧便携。这使得在单个 GPU 上进行微调变得可行，并且可以在推理时灵活地针对不同领域切换适配器。

本指南将介绍如何使用 diffusers 和 accelerate 库，以 LoRA 和 DoRA 对 Cosmos Predict 2.5 进行参数高效微调，并支持单 GPU 和多 GPU 训练。然后，我们将展示如何使用微调后的模型生成合成机器人轨迹，以用于下游的机器人学习 (https://www.nvidia.com/en-us/use-cases/robot-learning/) 任务。

https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#requirements 需求

• Python 3.10+
• PyTorch 2.5+ 及 CUDA
• diffusers（会自动拉取 transformers 和 peft）、accelerate
• 可选：安装 wandb 以监控训练
• 至少一块 80 GB 的 GPU 用于单 GPU 训练；推荐使用 8× H100 以加快迭代

在您的机器上安装依赖：

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

https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#preparing-data 数据准备

安装 diffusers 后，请导航至 examples/cosmos (https://github.com/terarachang/diffusers/tree/cosmos_predict_2.5_lora_clean/examples/cosmos) 查看示例代码。我们使用与 GR00T Dreams 后训练配方 (https://nvidia-cosmos.github.io/cosmos-cookbook/recipes/end2end/gr00t-dreams/post-training.html) 相同的数据集：

• 训练数据集 (https://huggingface.co/datasets/nvidia/GR1-100)：92 个机器人操作视频，并附有描述拾取和放置任务的文本提示。
• 测试数据集 (https://huggingface.co/datasets/nvidia/PhysicalAI-Robotics-GR00T-Eval)：50 个（提示，图像）对。模型应根据输入文本提示和初始帧图像生成视频。

使用 download_and_preprocess_datasets.sh (https://github.com/terarachang/diffusers/blob/cosmos_predict_2.5_lora_clean/examples/cosmos/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 └── ...

https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#training 训练

在本节中，我们将逐步介绍 train_cosmos_predict25_lora.py (https://github.com/terarachang/diffusers/blob/cosmos_predict_2.5_lora_clean/examples/cosmos/train_cosmos_predict25_lora.py) 中的实现。

https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#videodataset VideoDataset

VideoDataset 从 args.train_data_dir（在我们的示例中是 gr1_dataset/train）中将每个样本作为 (caption, video) 对加载。对于长度超过 args.num_frames 的视频，每个 epoch 会随机采样一个连续的 args.num_frames 窗口，从而实现时间增强。内部，VideoProcessor 来自 diffusers.video_processor，将原始帧调整大小并归一化为形状为 (channels, frames, height, width) 的张量。

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

https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#initialize-adapter 初始化适配器

Cosmos Predict 2.5 由三个子模块组成：

1. 一个将视频编码为潜变量的 VAE
2. 一个将文本提示编码为提示嵌入的文本编码器
3. 用于在潜空间中进行扩散的 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, )

冻结所有基础权重

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, # 设置为 True 以切换到 DoRA ) dit.add_adapter(lora_config) cast_training_params(dit, dtype=torch.float32) # LoRA 参数使用 fp32 ``

传递 use_dora=True 将切换到 DoRA，它会在应用低秩更新之前将每个权重分解为幅度和方向。训练循环中不需要其他更改。

https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#loss 损失函数

Cosmos Predict 2.5 使用整流流：模型被训练来预测将噪声样本线性传输到原始“干净”数据的速度。具体来说，在时间步 t，构建一个带噪声的插值 xt = σt·noise + (1−σt)·clean，其中 σt 是采样噪声水平，模型通过均方误差（MSE 损失）学习预测目标速度 noise − clean。视频的前两帧用作条件，因此它们的潜变量不会添加噪声。

训练损失遵循 Cosmos Predict 2.5 使用的整流流公式：

``

使用 logit-normal 分布采样时间步

sigma_t = sample_train_sigma_t(bsz, distribution=‘logitnormal’, device=device)

整流流在干净潜变量和噪声之间插值

xt = noise * sigma_t + clean_latent * (1 - sigma_t)

条件生成：DiT 以视频的前两帧、时间步和提示嵌入为条件

cond_indicator 和 cond_mask 对前两帧的值为 1，对其他帧的值为 0

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, )[0]

MSE 损失仅在非条件帧上计算

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()) ``

https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#optimizer-and-scheduler 优化器与调度器

我们使用 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, ) ``

https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#checkpointing 检查点

LoRA 权重以 diffusers 格式每 args.checkpointing_epochs 个 epoch 保存一次：

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 文件，该文件就是推理时传递给流水线的适配器文件。

https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#training-command 训练命令

使用提供的 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 小时。

https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#running-inference-with-your-lora 使用你的 LoRA 进行推理

训练完成后，使用 eval_cosmos_predict25_lora.py (https://github.com/terarachang/diffusers/blob/cosmos_predict_2.5_lora_clean/examples/cosmos/eval_cosmos_predict25_lora.py) 从评估数据集生成视频。该脚本从 gr1_dataset/test 中读取配对的 .png 和 .txt 文件，为每个图像生成一个视频，并将 .mp4 文件写入 --output_dir。

https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#imagedataset 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}

https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#loading-the-pipeline-and-loradora-weights 加载流水线和 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 分解带来的任何推理开销。

https://huggingface.co/blog/nvidia/cosmos-fine-tuning-for-robot-video-generation#generating-initial-latent-noise 生成初始潜变量噪声

为了确保可重复性，arch_invariant_rand 函数通过 NumPy (https://www.nvidia.com/en-us/glossary/numpy/) 生成初始潜变量噪声，使噪声不受 GPU 架构的影响。如果可重复性不是问题，用户可以不用该函数。