H100 8장으로 2.5시간 만에 구현한 Cosmos 2.5 Robot Video Fine-tuning

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

2026년 5월 18일12분advanced

AI 요약

Context

실제 로봇 궤적 데이터 수집의 고비용 및 저효율 문제 해결을 위해 World Model 기반의 합성 데이터 생성 필요성 증대. 2B 파라미터 규모의 Full Fine-tuning 시 발생하는 막대한 메모리 비용과 Catastrophic Forgetting 리스크가 주요 병목 지점으로 작용.

Technical Solution

VAE, Text Encoder, DiT 가중치를 Frozen 상태로 유지하여 기본 지식 보존 및 메모리 점유율 최적화
DiT 내 Attention Projection(to_q, to_k, to_v, to_out.0) 및 Feedforward Layer에 LoRA Adapter를 주입한 효율적 학습 구조 설계
Rectified Flow 기반의 Velocity Prediction 학습을 통해 Noise 샘플을 Clean 데이터로 선형 이동시키는 MSE Loss 적용
수치적 안정성 확보를 위해 BF16 Mixed Precision 환경에서 Trainable LoRA 파라미터만 Float32로 Upcast
Weight를 Magnitude와 Direction으로 분해하는 DoRA 구조를 통해 저차원 Rank에서의 학습 안정성 강화
Temporal Augmentation을 위해 Epoch마다 랜덤하게 연속 윈도우 프레임을 샘플링하는 VideoDataset 구현

실천 포인트

1. 메모리 제약이 극심한 경우 LoRA r=8부터 시작하고, 저차원 학습 시 불안정성이 관찰될 경우 DoRA r=32 도입 검토

2. 정밀한 Task 수행 능력이 필요하다면 Rank를 높여 Adapter의 Capacity를 확장

3. 비디오 생성 모델 학습 시 수치 안정성을 위해 LoRA 파라미터의 Float32 Upcast 적용 여부 확인

태그

#World Model #DORA #DiT #Rectified Flow #LoRA

원문 읽기