로봇 비디오 생성을 위한 LoRA 및 DoRA 기반 NVIDIA Cosmos Predict 2.5 미세 조정 가이드

NVIDIA Cosmos Predict 2.5는 텍스트, 이미지 또는 비디오 클립을 조건으로 물리적으로 타당한 비디오를 생성할 수 있는 대규모 월드 모델 (World Model)입니다. 이를 로봇 조작 (Robot Manipulation)이나 특정 카메라 시점과 같은 특정 도메인에 적응시키기 위해서는 여전히 타겟 미세 조정 (Fine-tuning)이 필요합니다.

로봇 정책 (Robot Policies)을 학습시키려면 시연 데이터 (Demonstration Data)가 필요하지만, 실제 로봇의 궤적 (Trajectories)을 수집하는 것은 느리고 비용이 많이 듭니다. 미세 조정된 비디오 월드 모델을 사용하여 합성 궤적 (Synthetic Trajectories)을 생성하는 것은 확장 가능한 대안을 제공합니다. 그러나 2B 파라미터 모델을 전체 미세 조정 (Full Fine-tuning)하는 것은 비용이 많이 들고 일반적인 지식을 잃어버리는 파괴적 망각 (Catastrophic Forgetting)의 위험이 있습니다. LoRA 및 DoRA는 동결된 베이스 모델 (Frozen Base Model)에 작은 학습 가능한 어댑터 (Adapter) 모듈을 주입하여, 어댑터 파일의 크기를 작고 휴대 가능하게 유지하면서 메모리 요구 사항을 줄여줍니다. 이를 통해 단일 GPU에서 미세 조정을 수행하는 것이 실용적이며, 추론 (Inference) 시 다양한 도메인에 맞춰 어댑터를 유연하게 교체할 수 있습니다.

이 가이드는 diffusers 및 accelerate 라이브러리를 사용하여 단일 및 멀티 GPU 학습을 모두 지원하는 LoRA 및 DoRA 기반의 Cosmos Predict 2.5 매개변수 효율적 미세 조정 (Parameter-Efficient Fine-Tuning) 과정을 안내합니다. 그런 다음 미세 조정된 모델을 사용하여 다운스트림 로봇 학습 (Downstream Robot Learning) 작업을 위한 합성 로봇 궤적을 생성하는 방법을 보여줍니다.

• Python 3.10+
• CUDA를 포함한 PyTorch 2.5+
• diffusers (transformers 및 peft를 자동으로 포함함), accelerate
• 선택 사항: 학습 모니터링을 위한 wandb 설치
• 단일 GPU 학습을 위해 최소 80 GB GPU 1개 필요; 빠른 반복을 위해 8× H100 권장

기기에 종속성을 설치하세요:

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

diffusers를 설치한 후, examples/cosmos로 이동하여 예제 코드를 탐색하세요.

우리는 GR00T Dreams 사후 학습 레시피와 동일한 데이터셋을 사용합니다:

• 학습 데이터셋 (Training Dataset): 집기 및 놓기 (Pick-and-place) 작업을 설명하는 텍스트 프롬프트가 포함된 92개의 로봇 조작 비디오.
• 테스트 데이터셋 (Test Dataset): 50개의 (프롬프트, 이미지) 쌍. 모델은 입력된 텍스트 프롬프트와 초기 프레임 이미지를 기반으로 비디오를 생성해야 합니다.

download_and_preprocess_datasets.sh를 사용하여 훈련 및 테스트 데이터셋을 다운로드하고 전처리합니다:

bash download_and_preprocess_datasets.sh

결과로 생성된 훈련 데이터셋(training dataset) 폴더의 구조는 다음과 같습니다:

gr1_dataset/train
├── metas/
│ └── *.txt
...

평가(eval) 데이터셋은 (프롬프트, 이미지) 쌍에 대한 .txt 및 .png 파일이 쌍을 이루는 평면 디렉토리(flat directory) 구조입니다:

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

이 섹션에서는 train_cosmos_predict25_lora.py의 구현 과정을 살펴봅니다.

VideoDataset은 args.train_data_dir(예시에서는 gr1_dataset/train)로부터 각 샘플을 (caption, video) 쌍으로 로드합니다. args.num_frames보다 긴 비디오의 경우, 매 에포크(epoch)마다 args.num_frames 크기의 임의의 연속된 윈도우(contiguous window)를 샘플링하여 시간적 증강 (temporal augmentation)을 가능하게 합니다. 내부적으로는 diffusers.video_processor의 VideoProcessor가 원본 프레임을 (channels, frames, height, width) 형태의 텐서(tensor)로 크기를 조정(resize)하고 정규화(normalize)합니다.

train_dataset = VideoDataset(
    dataset_dir=args.train_data_dir,
    num_frames=args.num_frames,
    ...
)

Cosmos Predict 2.5는 세 가지 서브모듈(submodules)로 구성됩니다:

• 비디오를 잠재 공간(latents)으로 인코딩하는 VAE
• 텍스트 프롬프트를 프롬프트 임베딩(prompt embeddings)으로 인코딩하는 텍스트 인코더 (text encoder)
• 잠재 공간에서 확산(diffusion)을 수행하는 DiT

훈련 중에는 모든 VAE, 텍스트 인코더, 그리고 DiT 가중치(weights)가 고정(frozen)됩니다. LoRA 어댑터(adapters)는 DiT의 어텐션 투영(attention projections) (to_q, to_k, to_v, to_out.0) 및 피드포워드 레이어 (feedforward layers) (ff.net.0.proj, ff.net.2)에 주입됩니다. 이후 학습 가능한 LoRA 파라미터(parameters)는 bf16 혼합 정밀도 (mixed precision) 환경에서 수치적 안정성을 위해 float32로 업캐스트(upcast)됩니다.

from diffusers import Cosmos2_5_PredictBasePipeline
from peft import LoraConfig

pipe = Cosmos2_5_PredictBasePipeline.from_pretrained(
    ...
)

use_dora=True를 전달하면 DoRA로 전환되며, DoRA는 저차원 업데이트 (low-rank update)를 적용하기 전에 각 가중치를 크기(magnitude)와 방향(direction)으로 분해합니다. 훈련 루프(training loop)에 필요한 다른 변경 사항은 없습니다.

Cosmos Predict 2.5는 **교정된 흐름 (rectified flow)**을 사용합니다. 모델은 노이즈 샘플을 원래의 "깨끗한 (clean)" 데이터로 선형적으로 이동시키는 속도 (velocity)를 예측하도록 훈련됩니다. 구체적으로, 타임스텝 (timestep) $t$에서, 샘플링된 노이즈 레벨 $\sigma_t$에 따라 노이즈가 섞인 보간 (noisy interpolation) xt = σt·noise + (1−σt)·clean이 구성되며, 모델은 평균 제곱 오차 (MSE loss)를 통해 목표 속도 noise − clean을 예측하는 법을 학습합니다. 비디오의 처음 두 프레임은 조건 (conditioning)으로 사용되므로, 해당 프레임의 잠재 변수 (latents)에는 노이즈가 추가되지 않습니다.

훈련 손실 (training loss)은 Cosmos Predict 2.5에서 사용되는 교정된 흐름 공식을 따릅니다:

# logit-normal 분포로 타임스텝 샘플링
sigma_t = sample_train_sigma_t(bsz, distribution='logitnormal', device=device)
# Rectified flow는 깨끗한 잠재 변수와 노이즈 사이를 보간함
...

최적화 도구 (optimizer)로는 torch.optim.AdamW를 사용하며, 스케줄러 (scheduler)로는 diffusers.optimization의 get_linear_schedule_with_warmup을 사용합니다. 스케줄러는 scheduler_warm_up_steps 동안 학습률 (learning rate)을 선형적으로 워밍업 (warm up)하여 scheduler_f_max × learning_rate에서 정점에 도달한 후, 남은 num_training_steps 동안 scheduler_f_min × learning_rate까지 선형적으로 감소(decay)시킵니다.

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

LoRA 가중치는 매 args.checkpointing_epochs 에포크 (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에 pytorch_lora_weights.safetensors 파일을 작성하며, 이 파일은 추론 (inference) 시점에 파이프라인 (pipeline)에 전달할 어댑터 (adapter) 파일이 됩니다.

제공된 쉘 스크립트 (shell script)를 시작점으로 사용하세요:

export MODEL_NAME="nvidia/Cosmos-Predict2.5-2B"
export DATA_DIR="gr1_dataset/train"
export OUT_DIR=YOUR_OUTPUT_DIR
...

** lora_rank**는 저차원 분해 (low-rank decomposition)의 랭크를 제어합니다. 랭크가 높을수록 학습 가능한 파라미터 (trainable parameters)가 많아지고 표현 능력 (expressive capacity)이 커지지만, 더 많은 메모리와 더 큰 어댑터 (adapter) 파일이 필요합니다. 여기서는 rank=32를 시작점으로 사용하며, 그 결과 약 50M 개의 학습 가능한 파라미터가 생성됩니다.

** lora_alpha**는 LoRA 업데이트에 적용되는 스케일링 인자 (scaling factor)입니다. 가중치 델타 (weight delta)는 동결된 베이스 가중치 (frozen base weights)에 더해지기 전에 다음과 같이 스케일링됩니다:

lora_alpha / lora_rank

(여기서 수행한 것처럼) lora_alpha = lora_rank로 설정하면 이 스케일 인자가 1.0으로 유지되어, LoRA 업데이트가 추가적인 감쇄 (dampening) 없이 전체 강도로 적용됩니다. LoRA 대신 DoRA를 사용하려면 명령에 --use_dora를 추가하십시오.

멀티 GPU 학습의 경우, accelerate가 분산을 자동으로 처리합니다. 경험적으로, 이 작업에서는 100 에포크 (epochs) 동안 학습하는 것만으로도 이미 괜찮은 결과를 얻을 수 있음을 확인했습니다. 이는 단일 H100에서 17시간, 8개의 H100 GPU에서는 2.5시간이 소요됩니다.

학습이 완료되면 eval_cosmos_predict25_lora.py를 사용하여 평가 데이터셋 (eval dataset)으로부터 비디오를 생성합니다. 이 스크립트는 gr1_dataset/test에서 쌍을 이루는 .png 및 .txt 파일을 읽어 각각에 대한 비디오를 생성하고, .mp4 파일을 --output_dir에 작성합니다.

ImageDataset은 .txt 파일을 프롬프트 문자열 (prompt string)로 읽어들이고, 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)
...

from diffusers import Cosmos2_5_PredictBasePipeline
pipe = Cosmos2_5_PredictBasePipeline.from_pretrained(
"nvidia/Cosmos-Predict2.5-2B",
...

fuse_lora는 어댑터 가중치를 베이스 모델에 병합하여, LoRA/DoRA 분해로 인한 추론 오버헤드 (inference overhead)를 제거합니다.

재현성 (reproducibility)을 보장하기 위해, arch_invariant_rand 함수는 NumPy를 통해 초기 잠재 노이즈 (latent noise)를 생성하므로 노이즈가 GPU 아키텍처에 영향을 받지 않습니다 (invariant). 재현성이 중요하지 않다면, 사용자는 파이프라인 (pipeline)에 입력 노이즈를 제공할 필요가 없습니다.

# 잠재 변수 (latent)와 동일한 형상을 가진 무작위 노이즈로부터 생성 시작
latent_shape = pipe.get_latent_shape_cthw(args.height, args.width, args.num_output_frames)
noises = arch_invariant_rand(
...

export LORA_DIR=YOUR_ADAPTER_DIR
export DATA_DIR="gr1_dataset/test"
export OUT_DIR=YOUR_EVAL_OUTPUT_DIR
...

LoRA 없이 베이스 모델 (base model)을 평가하려면 --lora_dir을 생략하십시오.

Sampson Error는 매칭된 키포인트 (keypoints)로부터 그에 대응하는 에피폴라 라인 (epipolar lines)까지의 거리를 측정하는 기하학적 오차 지표 (geometric error metric)입니다. 생성된 비디오의 맥락에서, 낮은 Sampson error는 프레임 간(또는 카메라 뷰 간)의 움직임이 기하학적으로 일관됨을 의미합니다. 값이 높을수록 지터 (jitter), 환각된 움직임 (hallucinated motion), 또는 다중 뷰 불일치 (multi-view inconsistencies)를 나타냅니다.

우리는 Cosmos Predict 평가 가이드를 따르며, 두 가지 지표를 사용하여 생성된 비디오의 기하학적 품질을 평가합니다:

시간적 Sampson Error (Temporal Sampson Error): 단일 카메라 뷰 내의 연속된 프레임 사이에서 계산되며, 시간적 안정성 (temporal stability)을 측정합니다.
교차 뷰 Sampson Error (Cross-view Sampson Error): 서로 다른 카메라 뷰의 동시 프레임 사이에서 계산되며, 다중 뷰 기하학적 정렬 (multi-view geometric alignment)을 측정합니다.

우리는 Cosmos Reason2를 LLM 심사위원 (LLM judge)으로 사용하여 각 예시를 1점에서 5점 사이로 점수를 매깁니다. 우리는 두 가지 루브릭 (rubrics)을 설계했습니다:

물리적 타당성 (Physical plausibility) (video_physics.yaml): 심사위원은 텍스트 프롬프트 (text prompt)를 보지 않고, 비디오가 물리적 상식 (physical commonsense)을 따르는지 평가합니다.
지시 이행 (Instruction following) (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.
...

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

테스트 세트의 처음 두 가지 예시에 대해 베이스 모델 (미세 조정 전), LoRA, 그리고 DoRA로 생성된 비디오를 비교합니다.

프롬프트 (Prompt): 왼손을 사용하여 원형 회색 매트 위에 있는 진한 녹색 오이를 베이지색 그릇 위로 집어 올리세요.

학습 전 (Before Training)	LoRA r=32	DoRA r=32

프롬프트 (Prompt): 오른손을 사용하여 분홍색 접시 중앙에 있는 오렌지 주스 팩을 초록색 그릇 중앙으로 집어 올리세요.

학습 전 (Before Training)	LoRA r=32	DoRA r=32

미세 조정 (Fine-tuning) 전에는 베이스 모델이 여러 측면에서 어려움을 겪습니다. 로봇 손이 분포 외 데이터 (Out-of-distribution)에 해당하여 모델이 이후 프레임에서 사람의 손을 환각 (Hallucinate)하게 만들고, 프롬프트에 지정된 올바른 손을 안정적으로 사용하지 못하며, 생성된 비디오에서 눈에 띄는 지터 (Jitter) 현상이 나타납니다. LoRA와 DoRA를 이용한 미세 조정은 이 세 가지 문제를 모두 해결합니다.

우리는 서로 다른 설정 하에 네 가지 어댑터 (Adapter)를 미세 조정합니다: 랭크 (Rank) 8과 32를 적용한 LoRA 및 DoRA입니다. 각 테스트 예시에 대해 서로 다른 시드 (Seed)를 사용하여 5개의 비디오를 생성하고, '평가 지표 (Evaluation Metrics)' 섹션에서 소개된 세 가지 지표를 사용하여 시드 전체의 평균 점수를 보고합니다.

결론 (Conclusion): 100 에포크 (Epochs, 8× H100에서 약 2.5시간) 동안의 학습만으로도 세 가지 지표 모두를 실질적으로 개선하기에 충분합니다. LoRA와 DoRA 모두 유사한 성능으로 수렴하며, 이는 DoRA의 추가적인 크기-방향 분해 (Magnitude-direction decomposition)가 성능을 저해하지 않으며 매우 낮은 랭크에서는 도움이 될 수 있지만, 여기서는 필수적이지 않음을 확인시켜 줍니다.

더 큰 랭크 (32 대 8)는 지시 이행 (Instruction following) 능력을 향상시키지만 (모델이 어떤 손을 사용하고 어떤 객체와 상호작용해야 하는지를 더 정밀하게 학습할 수 있는 용량을 갖게 됨), 기하학적 일관성 (Geometric consistency)이나 물리적 타당성 (Physical plausibility)을 개선하지는 않습니다. 우리는 이것이 기하학적 및 물리적 사전 지식 (Priors)이 월드 모델 (World model)의 동결된 가중치 (Frozen weights)에 의해 대부분 포착되어 있기 때문이라고 가설을 세웁니다. LoRA 어댑터는 분포를 도메인 내 (In-domain) 로봇 외형과 작업 구조로 이동시키기만 하면 되며, 이는 랭크 8에서도 달성 가능합니다.

DoRA와 LoRA 중 언제 무엇을 사용할 것인가: 메모리가 매우 부족하거나 어댑터 파일 크기가 중요한 경우, LoRA r=8로 시작하십시오. 예산이 충분하고 낮은 랭크(low rank)에서 LoRA 사용 시 학습 불안정성(training instability)이 관찰된다면, DoRA r=32가 합리적인 대안이 될 수 있습니다. 크기-방향 분해(magnitude–direction decomposition)가 학습을 안정화하는 데 도움이 될 수 있기 때문입니다.

Cosmos WFMs를 구축, 적응 및 배포하기 위한 단계별 워크플로우, 기술 레시피 및 구체적인 예시를 보려면 Cosmos Cookbook을 방문하십시오. Hugging Face와 GitHub에서 새로운 오픈 Cosmos 모델과 데이터셋을 탐색하거나 build.nvidia.com에서 모델을 테스트해 보세요. Cosmos Discord 채널에 참여하여 커뮤니티의 일원이 되어보세요. 이미 Cosmos를 사용 중이신가요? 기여하는 방법에 대해 더 자세히 알아보세요.