‏فاین‌تیون Cosmos Predict 2.5 با LoRA برای تولید ویدیوی ربات

NVIDIA Cosmos Predict 2.5 یه مدل جهانی (World Model) در مقیاس بزرگه که می‌تونه ویدیوهای فیزیکی‌باور بسازه؛ مشروط به متن، تصویر اولیه، یا کلیپ ویدیویی. یکی از کاربردهای مهمش اینه که به‌جای جمع‌آوری دیتای واقعی از ربات‌ها — که کند و گرونه — بشه مسیرهای مصنوعی تولید کرد و ازشون برای آموزش مدل‌های رباتیک استفاده کرد.

فاین‌تیون کامل یه مدل ۲ میلیارد پارامتری هم پرهزینه‌ست و هم ریسک «فراموشی فاجعه‌بار» داره — یعنی مدل دانش عمومیش رو از دست می‌ده. LoRA (Low-Rank Adaptation) و DoRA راه‌حل هستن: ماژول‌های کوچک و قابل‌آموزشی رو به مدل منجمد اضافه می‌کنن. نتیجه اینه که فقط ~۵۰ میلیون پارامتر آموزش می‌بینن، فایل آداپتور کوچک می‌مونه، و می‌شه روی یه GPU تکی هم کار کرد.

معماری مدل از سه بخش تشکیل شده: یه VAE که ویدیو رو به فضای پنهان (latent) تبدیل می‌کنه، یه text encoder که متن رو encode می‌کنه، و یه DiT (Diffusion Transformer) که در فضای latent کار می‌کنه. در طول آموزش، وزن‌های هر سه بخش منجمد می‌مونن و آداپتورهای LoRA فقط به لایه‌های attention و feedforward مربوط به DiT تزریق می‌شن:

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)

برای تعویض به DoRA کافیه use_dora=True بذاری — هیچ تغییر دیگه‌ای در حلقه‌ی آموزش لازم نیست. DoRA هر وزن رو به مؤلفه‌ی بزرگی (magnitude) و جهت (direction) تجزیه می‌کنه و بعد آپدیت low-rank رو اعمال می‌کنه.

تابع loss بر اساس Rectified Flow کار می‌کنه: مدل یاد می‌گیره سرعتی رو پیش‌بینی کنه که یه نمونه‌ی نویزی رو به داده‌ی تمیز منتقل می‌کنه. دو فریم اول هر ویدیو به‌عنوان شرط conditioning استفاده می‌شن و نویزی بهشون اضافه نمی‌شه. بقیه‌ی فریم‌ها با MSE loss آموزش می‌بینن.

برای inference، آداپتور LoRA رو لود می‌کنی و با fuse_lora ادغامش می‌کنی تا overhead اضافی نداشته باشی:

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)

ارزیابی مدل با دو روش انجام می‌شه: Sampson Error که سازگاری هندسی بین فریم‌ها رو می‌سنجه (مقدار کمتر یعنی حرکت طبیعی‌تر)، و LLM-as-a-Judge که از مدل Cosmos Reason2 استفاده می‌کنه تا هم فیزیک‌باوری ویدیو و هم پیروی از دستور متنی رو از ۱ تا ۵ امتیاز بده. آموزش با ۱۰۰ epoch روی ۸ عدد H100 حدود ۲.۵ ساعت طول می‌کشه.

نکات کلیدی:

• Cosmos Predict 2.5 یه مدل ۲B پارامتری برای تولید ویدیوی فیزیکی‌باوره
• LoRA/DoRA فقط ~۵۰M پارامتر آموزش می‌دن و فاین‌تیون روی یه GPU تکی ممکنه
• rank=32 و lora_alpha=32 به‌عنوان نقطه‌ی شروع پیشنهاد می‌شه
• DoRA با اضافه کردن --use_dora فعال می‌شه، بدون نیاز به تغییر بقیه‌ی کد
• ارزیابی با Sampson Error (هندسی) و LLM Judge (کیفی) انجام می‌شه
• هدف نهایی تولید داده‌ی مصنوعی برای آموزش سیاست‌های رباتیک است