NVIDIA视频生成微调：省2小时+适配率翻倍

新媒网跨境了解到，NVIDIA Cosmos Predict 2.5，这款先进的视频生成模型，专注于用文本、图片或视频片段生成物理真实的视频。当应用到特定领域，比如机器人操作场景或者某些特定的镜头角度时，模型需要进行进一步的微调。然而，对于跨境从业者来说，这种模型动辄几十亿参数的完整微调不仅耗时耗力，而且可能导致遗忘掉模型的通用知识。

基于此，技术界研发了LoRA（低秩适配器）和DoRA（动态优化适配器），这两种轻量化模块通过“插入式”的适配方式可以帮我们实现小规模、高效的模型调整。这种方法能显著减少显存需求，并且让适配器文件保持小巧而便于共享，即使只有单卡GPU也能完成训练。同时，这些适配器可以在不同应用领域间灵活切换。对于国内跨境领域的合作伙伴来说，这无疑是一个优化资源、提升效率的好办法。下面新媒网跨境将结合实战案例，逐步拆解如何为机器人操作生成视频并进行微调。

一、基础准备工作

首先，你需要准备以下环境与设备：

Python 3.10或更高版本。
PyTorch 2.5以上，并确保支持CUDA。
相关Python库：diffusers（包含transformers和peft）、accelerate。
可选工具：安装wandb，方便监控训练过程。
设备推荐：单张80GB GPU进行单机微调，或者集群使用8张H100提升速度。

安装上述工具的命令可在终端执行：

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

二、数据准备

进入examples/cosmos文件夹，你会发现示例代码以及训练所需数据的准备脚本。以下是数据结构的主要说明：

训练数据集

包含92段机器人操作视频及与之关联的文字提示，用于描述例如“抓取—放置”等任务。

数据结构如下：

gr1_dataset/train
├── metas/
│   └── *.txt
├── videos/
│   └── *.mp4
└── metadata.csv

测试数据集

包含50组“文字提示+图像帧”对，用于验证模型生成的视频质量。

数据结构如下：

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

运行以下命令下载并处理数据：

bash download_and_preprocess_datasets.sh

三、模型微调步骤解析

1. 构建视频数据集加载机制

我们定义一个VideoDataset类，能够将训练集加载为“文字描述+视频数据”的配对。在训练过程中，我们使用随机取样技术让模型每轮学习不同的时间窗口，以充分利用视频中的动态细节。

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

2. 初始化LoRA适配器

NVIDIA Cosmos Predict 2.5由三个主要组件构成：

VAE模块：将视频编码为潜在变量。
文本编码器：处理文字并生成文本嵌入。
DiT（扩散模型）：在潜在空间中生成实际视频。

在微调过程中，我们将VAE、文本编码器以及DiT的主要权重冻结，同时采用LoRA插入到DiT的注意力投影和前馈层中。

以下代码展示了如何为DiT模块添加LoRA适配器：

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参数
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,  # 根据需求切换至DoRA
)

# 插入适配器并启用训练
dit.add_adapter(lora_config)
cast_training_params(dit, dtype=torch.float32)  # 将LoRA参数设定为fp32

3. 损失函数设计

我们采用预测“速度梯度”的方式训练模型，通过均方误差（MSE）损失评估生成视频的精度。

pred_velocity = dit(xt, cond_mask, in_timestep, prompt_embeds, padding_mask, return_dict=False)[0]
target_velocity = noise - clean_latent
loss = F.mse_loss(pred_velocity.float(), target_velocity.float())

4. 优化器和学习率调节

使用AdamW作为优化器，并结合线性增温学习率调控器实现长时间稳定微调。

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,
)

5. 模型保存与加载

微调过程中，每隔一定轮数保存当前的参数状态，以便后续使用或排查。

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)

6. 实操训练指令

将所需参数和路径设置为环境变量后，通过以下命令启动训练：

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 \
--height 432 --width 768 \
--gradient_checkpointing --allow_tf32 \
--report_to=wandb \
--lora_rank 32