Giới thiệu: Nhu cầu tạo Video Robot chân thực
NVIDIA Cosmos Predict 2.5, một "mô hình thế giới" quy mô lớn, hứa hẹn khả năng tạo video chân thực về mặt vật lý dựa trên văn bản, hình ảnh hoặc các đoạn video khác. Tuy nhiên, để thực sự ứng dụng vào các lĩnh vực cụ thể như thao tác robot hay góc nhìn camera chuyên biệt, việc tinh chỉnh mục tiêu là không thể thiếu. 🤖
Việc huấn luyện các chính sách robot đòi hỏi dữ liệu minh họa thực tế, nhưng việc thu thập quỹ đạo từ robot thật lại rất tốn kém và mất thời gian. Giải pháp tiềm năng là tạo ra các quỹ đạo tổng hợp bằng mô hình thế giới video đã được tinh chỉnh. Tuy nhiên, việc tinh chỉnh toàn bộ một mô hình với 2 tỷ tham số không chỉ đắt đỏ mà còn đối mặt với rủi ro "quên thảm khốc" (catastrophic forgetting) những kiến thức tổng quát đã học. Đây là một thách thức lớn! 😥
LoRA/DoRA: Giải pháp tinh chỉnh hiệu quả cho Cosmos Predict 2.5
Để giải quyết vấn đề này, các kỹ thuật như LoRA (Low-Rank Adaptation) và DoRA (Damped LoRA) đã xuất hiện như những "vị cứu tinh". Chúng cho phép chúng ta tích hợp các module adapter nhỏ, có thể huấn luyện được, vào mô hình nền đã bị đóng băng. Điều này giúp giảm đáng kể yêu cầu về bộ nhớ, đồng thời giữ cho các tệp adapter nhỏ gọn và dễ di chuyển. Nhờ đó, việc tinh chỉnh trở nên khả thi trên một GPU đơn lẻ và cho phép linh hoạt hoán đổi adapter cho các miền khác nhau trong quá trình suy luận. Thật tiện lợi! 💡
Bài viết này của Hugging Face và NVIDIA cung cấp hướng dẫn chi tiết về cách thực hiện tinh chỉnh hiệu quả tham số (parameter-efficient fine-tuning) cho Cosmos Predict 2.5 bằng LoRA và DoRA, sử dụng thư viện diffusers và accelerate. Chúng ta cũng sẽ tìm hiểu cách sử dụng mô hình đã tinh chỉnh để tạo ra các quỹ đạo robot tổng hợp cho các tác vụ học máy robot sau này.
Yêu cầu và Chuẩn bị Dữ liệu
Để bắt đầu, bạn cần đảm bảo hệ thống của mình đáp ứng các yêu cầu sau: * Python 3.10+ * PyTorch 2.5+ với CUDA * Các thư viện diffusers (sẽ tự động kéo theo transformers và peft), accelerate * Tùy chọn: wandb để giám sát quá trình huấn luyện
Về phần cứng, bạn cần tối thiểu một GPU 80 GB cho việc huấn luyện đơn GPU; tuy nhiên, 8 GPU H100 được khuyến nghị để tăng tốc độ lặp lại.
Sau khi cài đặt diffusers, bạn có thể truy cập mã ví dụ tại examples/cosmos trên GitHub. Bộ dữ liệu được sử dụng tương tự như công thức hậu huấn luyện GR00T Dreams: * Tập huấn luyện: nvidia/GR1-100 trên Hugging Face gồm 92 video thao tác robot với các gợi ý văn bản mô tả tác vụ "nhặt và đặt". * Tập kiểm tra: nvidia/PhysicalAI-Robotics-GR00T-Eval gồm 50 cặp (gợi ý, hình ảnh).
Bạn có thể tải và tiền xử lý dữ liệu bằng script download_and_preprocess_datasets.sh. Cấu trúc thư mục dữ liệu sau khi xử lý sẽ rất rõ ràng, giúp quản lý dễ dàng.
Quy trình Tinh chỉnh Đơn giản nhưng Mạnh mẽ
Quá trình tinh chỉnh được triển khai trong train_cosmos_predict25_lora.py. Dưới đây là những điểm đáng chú ý:
* VideoDataset: Tải mỗi mẫu dưới dạng cặp (caption, video). Đối với các video dài hơn số khung hình mong muốn, nó sẽ lấy ngẫu nhiên một cửa sổ liên tục các khung hình trong mỗi epoch, cho phép tăng cường dữ liệu theo thời gian. * Tối ưu hóa với LoRA/DoRA: Cosmos Predict 2.5 bao gồm ba submodule chính: VAE, bộ mã hóa văn bản và DiT. Trong quá trình huấn luyện, tất cả các trọng số của VAE, bộ mã hóa văn bản và DiT đều được đóng băng. Các adapter LoRA được chèn vào các phép chiếu chú ý (to_q, to_k, to_v, to_out.0) và các lớp feedforward (ff.net.0.proj, ff.net.2) của DiT. Việc sử dụng use_dora=True sẽ chuyển sang DoRA mà không cần thay đổi thêm trong vòng lặp huấn luyện. Điều này chứng tỏ sự linh hoạt của framework! * Hàm mất mát và Bộ tối ưu hóa: Mô hình sử dụng rectified flow, được huấn luyện để dự đoán vận tốc chuyển đổi tuyến tính từ mẫu nhiễu sang dữ liệu "sạch" gốc. Hàm mất mát MSE được tính chỉ trên các khung hình không điều kiện. Chúng ta dùng torch.optim.AdamW làm bộ tối ưu hóa và get_linear_schedule_with_warmup để điều chỉnh tốc độ học. * Lưu Checkpoint: Trọng số LoRA được lưu theo định dạng diffusers sau mỗi vài epoch, tạo ra tệp pytorch_lora_weights.safetensors dùng để suy luận.
Lệnh huấn luyện thực tế: Để khởi động quá trình, bạn có thể dùng lệnh sau làm điểm bắt đầu: ```bash 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 kiểm soát rank của phân tách hạng thấp (rank=32 cho khoảng 50 triệu tham số huấn luyện). lora_alpha là một yếu tố tỷ lệ; đặt lora_alpha = lora_rank giữ tỷ lệ này ở 1.0, áp dụng bản cập nhật LoRA đầy đủ. Đối với huấn luyện đa GPU, accelerate tự động xử lý phân phối. Đáng chú ý, huấn luyện 100 epoch chỉ mất 17 giờ trên một H100 và chỉ 2.5 giờ trên 8 GPU H100. Một sự tiết kiệm thời gian đáng kể! ⏰
Tạo Video Robot với Mô hình đã Tinh chỉnh
Sau khi hoàn tất huấn luyện, bạn có thể sử dụng eval_cosmos_predict25_lora.py để tạo video từ tập dữ liệu đánh giá. Script sẽ đọc các cặp tệp .png và .txt và xuất ra các tệp .mp4.
* ImageDataset: Đọc tệp .txt thành chuỗi prompt và tải .png dưới dạng PIL.Image.Image. * Tải Pipeline và Trọng số LoRA/DoRA: Rất đơn giản, chỉ cần tải pipeline cơ sở và sau đó dùng pipe.load_lora_weights("/path/to/lora/checkpoint") rồi pipe.fuse_lora(). Chức năng fuse_lora sẽ hợp nhất trọng số adapter vào mô hình cơ sở, loại bỏ mọi chi phí suy luận phát sinh từ việc phân tách LoRA/DoRA. * Tạo nhiễu latent ban đầu: Để đảm bảo khả năng tái tạo, hàm arch_invariant_rand tạo nhiễu latent ban đầu thông qua NumPy, làm cho nhiễu bất biến với kiến trúc GPU.
Lệnh suy luận: ```bash 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 ``` Để đánh giá mô hình cơ sở (base model) mà không có LoRA, bạn chỉ cần bỏ qua đối số --lora_dir.
Đánh giá Hiệu suất: Chất lượng và Độ chính xác
Các nhà nghiên cứu đã đánh giá chất lượng hình học của video được tạo ra bằng hai chỉ số Sampson Error và sử dụng Cosmos Reason2 làm LLM-as-a-Judge.
* Sampson Error: Chỉ số lỗi hình học này đo khoảng cách từ các điểm khóa khớp đến đường cực của chúng. Sampson Error thấp cho thấy chuyển động giữa các khung hình nhất quán về mặt hình học, tránh hiện tượng giật, chuyển động ảo hoặc không nhất quán đa khung hình. * Temporal Sampson Error: Tính giữa các khung hình liên tiếp trong cùng một góc nhìn. * Cross-view Sampson Error: Tính giữa các khung hình đồng thời từ các góc nhìn khác nhau. * LLM-as-a-Judge (Cosmos Reason2): Mô hình ngôn ngữ lớn này chấm điểm từ 1 đến 5 dựa trên hai tiêu chí: * Tính chân thực vật lý: Đánh giá xem video có tuân thủ các quy tắc vật lý thông thường không, không cần nhìn vào prompt văn bản. * Khả năng làm theo hướng dẫn: Đánh giá xem video có hoàn thành đúng tác vụ được mô tả trong prompt hay không.
Kết quả: Phân tích Định tính và Định lượng * Phân tích Định tính: Trước khi tinh chỉnh, mô hình cơ sở gặp nhiều khó khăn: tay robot bị sai lệch phân phối (out-of-distribution) dẫn đến việc mô hình "ảo giác" ra tay người; không sử dụng đúng tay robot được chỉ định trong prompt; và video tạo ra bị giật đáng kể. Việc tinh chỉnh với LoRA và DoRA đã khắc phục hiệu quả cả ba vấn đề này! Các ví dụ minh họa cho thấy sự cải thiện rõ rệt, đặc biệt là trong việc tạo ra chuyển động robot tự nhiên và tuân thủ hướng dẫn. * Phân tích Định lượng: Nghiên cứu đã tinh chỉnh bốn adapter (LoRA và DoRA với rank 8 và 32). Kết quả cho thấy: * Sampson Error (càng thấp càng tốt): Cả lỗi Temporal và Cross-view Sampson đều giảm đáng kể sau khi tinh chỉnh, cho thấy sự ổn định theo thời gian và tính nhất quán hình học đa khung hình được cải thiện. * Điểm chân thực vật lý (càng cao càng tốt): Các mô hình đã tinh chỉnh tạo ra video tuân thủ tốt hơn các quy tắc vật lý thông thường so với mô hình cơ sở. * Điểm làm theo hướng dẫn (càng cao càng tốt): Các mô hình đã tinh chỉnh hoàn thành tác vụ mô tả trong prompt một cách đáng tin cậy hơn, bao gồm cả việc sử dụng đúng tay và tương tác với các đối tượng cụ thể.
Kết luận quan trọng: Chỉ cần 100 epoch huấn luyện (khoảng 2.5 giờ trên 8 GPU H100) là đủ để cải thiện đáng kể cả ba chỉ số. Cả LoRA và DoRA đều đạt hiệu suất tương tự. Rank lớn hơn (32 so với 8) giúp tăng khả năng làm theo hướng dẫn, nhưng không cải thiện tính nhất quán hình học hoặc tính chân thực vật lý. Điều này gợi ý rằng các kiến thức vật lý và hình học cơ bản đã được mô hình thế giới đóng băng nắm bắt tốt; adapter LoRA chỉ cần điều chỉnh phân phối hướng tới giao diện robot và cấu trúc tác vụ cụ thể của miền, điều này có thể đạt được ngay cả với rank 8.
Khi nào nên dùng DoRA hay LoRA? Nếu bộ nhớ hạn chế hoặc kích thước tệp adapter là yếu tố quan trọng, hãy bắt đầu với LoRA r=8. Nếu bạn có ngân sách và nhận thấy sự bất ổn trong huấn luyện với LoRA ở rank thấp, DoRA r=32 là một lựa chọn thay thế hợp lý, vì phân tách độ lớn-hướng có thể giúp ổn định quá trình học.
Lời Kết và Tài nguyên Tham khảo
Việc tinh chỉnh NVIDIA Cosmos Predict 2.5 với LoRA/DoRA mở ra một hướng đi đầy hứa hẹn để tạo ra dữ liệu tổng hợp chất lượng cao, thúc đẩy nghiên cứu và phát triển trong lĩnh vực học máy robot. Đây là một bước tiến quan trọng giúp giảm thiểu chi phí và thời gian thu thập dữ liệu thực tế, đồng thời tăng cường khả năng của robot trong các tác vụ phức tạp.
Để tìm hiểu sâu hơn về Cosmos, bạn có thể khám phá: * Cosmos Cookbook để biết các quy trình làm việc và ví dụ chi tiết. * Các mô hình và bộ dữ liệu Cosmos mở trên Hugging Face và GitHub. * Thử nghiệm mô hình trên build.nvidia.com. * Tham gia cộng đồng trên kênh Discord của Cosmos.