Trender — AI 트렌드 자동 수집

Image-Text-to-Text • Updated • 209k • 77

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.

Python 3.10+
PyTorch 2.5+ with CUDA
```
diffusers
```
(pulls in
```
transformers
```
and
```
peft
```
automatically),
```
accelerate
```
Optional: install
```
wandb
```
to 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

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
└── ...

In this section, we walk through the implementation in train_cosmos_predict25_lora.py.

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

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.

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

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

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.

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.

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

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}

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.

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)

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

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.

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?

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.

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.

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 ourCosmos Cookbookfor step-by-step workflows, technical recipes, and concrete examples for building, adapting, and deploying Cosmos WFMs.
Explore new open Cosmos models and datasets onHugging FaceandGitHubor try models onbuild.nvidia.com.
Be part of the community and join ourCosmos Discord channel.
Already using Cosmos? Learn more abouthow to contribute.

Image-Text-to-Text • Updated • 209k • 77

Viewer • Updated • 126 • 269 • 7

Fine-Tuning NVIDIA Cosmos Predict 2.5 with LoRA/DoRA for Robot Video Generation

추출된 키워드

원문