使用LoRA/DoRA微調NVIDIA Cosmos Predict 2.5以生成機器人影片
本文介紹瞭如何透過引數高效微調技術LoRA和DoRA，在單個GPU上微調NVIDIA Cosmos Predict 2.5世界模型，生成用於機器人學習的合成影片軌跡。文章詳細說明了資料處理、介面卡初始化、訓練迴圈、推理方法及評估指標。
文章情報
要點
- LoRA和DoRA允許在凍結基模型的情況下，透過小型可訓練介面卡進行高效微調，避免災難性遺忘並降低記憶體需求。
- 訓練使用92個機器人操作影片資料集，透過rectified flow損失函式和MSE損失進行最佳化。
- 微調後的模型可生成合成影片軌跡，用於下游機器人策略學習。
為什麼重要
這條新聞值得關注，因為LoRA和DoRA允許在凍結基模型的情況下，透過小型可訓練介面卡進行高效微調，避免災難性遺忘並降低記憶體需求。
技術影響
可能影響模型選型、推理成本、產品能力和評測基準。
NVIDIA Cosmos Predict 2.5是一個大規模世界模型，能夠根據文本、影像或影片片段生成物理合理的影片。為了將其適應特定領域（如機器人操作或特定攝像機視角），需要針對性的微調。然而，對20億引數模型進行全引數微調成本高昂，且存在災難性遺忘的風險。LoRA（低秩適配）和DoRA（方向性低秩適配）透過在凍結基模型中注入小型可訓練介面卡模組，顯著降低記憶體需求，並使介面卡檔案保持小巧便攜，從而可以在單個GPU上完成微調，並在推理時靈活切換不同領域的介面卡。
首先，環境要求包括Python 3.10+、PyTorch 2.5+、diffusers和accelerate庫，並推薦使用至少一塊80GB視訊記憶體的GPU。資料準備方面，使用了GR00T Dreams後訓練配方中的資料集：訓練集包含92個機器人操作影片及對應的文本提示，測試集包含50個（提示、影像）對。透過預處理指令碼下載和整理資料。
訓練實現的核心是VideoDataset類，它載入（標題、影片）對，並對長影片進行隨機視窗取樣以實現時間增強。介面卡初始化時，將VAE、文本編碼器和DiT的所有權重凍結，僅對DiT的注意力投影層和前饋網路層注入LoRA介面卡。LoRA引數被上轉換為float32以保持數值穩定性。損失函式採用rectified flow形式：模型預測將噪聲線性傳輸到乾淨資料的速度，僅對非條件幀計算MSE損失。最佳化器使用AdamW，搭配帶線性預熱和衰減的學習率排程器。
訓練命令示例中，使用lora_rank=32（約5000萬可訓練引數），透過accelerate實現單機多卡訓練。實驗表明，100個epoch可在單個H100上約17小時完成，或在8塊H100上2.5小時完成。
推理時，載入微調後的介面卡並融合到基模型中，消除推理開銷。透過arch_invariant_rand函式生成初始潛變數噪聲，保證跨GPU架構的可重複性。評價指標包括Sampson誤差，用於衡量生成影片的幾何一致性。文章還提到了未來對TPA（時間位置注意力）和Kronecker介面卡的探索，以進一步提升效率。