Hướng dẫn sử dụng Cosmos của NVIDIA: Xây dựng mô hình thế giới thu nhỏ trên Google Colab

Trong bài hướng dẫn này, chúng ta sẽ khám phá cosmos-framework của NVIDIA từ góc độ thực tế, thân thiện với Colab, đồng thời nhìn nhận thẳng thắn về những giới hạn phần cứng khi chạy các checkpoint Cosmos 3 thực tế. Chúng ta bắt đầu bằng việc kiểm tra môi trường runtime hiện tại, khả năng của GPU, tính khả dụng của CUDA, bộ nhớ và dung lượng ổ đĩa để hiểu lý do tại sao việc suy luận (inference) toàn bộ Cosmos 3 là không khả thi trên phần cứng Colab tiêu chuẩn. Thay vì dừng lại ở đó, chúng ta sử dụng cấu trúc thực tế, giao diện CLI, lược đồ đầu vào (input schema) và các chế độ mô hình của framework làm nền tảng để triển khai một phiên bản thu nhỏ. Sau đó, chúng ta xây dựng và huấn luyện một mô hình thế giới Mixture-of-Transformers đa phương thức (omnimodal) nhỏ gọn, phản ánh ý tưởng cốt lõi của Cosmos: cơ chế chú ý (attention) đa phương thức được chia sẻ kết hợp với định tuyến chuyên gia (expert routing) dành riêng cho từng phương thức đối với các luồng văn bản, hình ảnh và hành động. Sử dụng dữ liệu thế giới vật lý tổng hợp, theo dõi hàm mất mát (training-loss) và cơ chế autoregressive rollout, chúng ta sẽ thấy cách mô hình học các mối quan hệ giữa các phương thức và dự đoán các trạng thái tiềm ẩn (latent states) trong tương lai theo một cách đơn giản nhưng có ý nghĩa về mặt kỹ thuật.

Thăm dò giới hạn phần cứng Colab

Chúng ta bắt đầu bằng việc chuẩn bị các tiện ích runtime và kiểm tra xem máy hiện tại có thể hỗ trợ suy luận Cosmos 3 một cách thực tế hay không. Chúng ta kiểm tra Python, PyTorch, CUDA, bộ nhớ GPU, khả năng tính toán và dung lượng ổ đĩa khả dụng để so sánh môi trường của chúng ta với các yêu cầu phần cứng thực tế. Sau đó, chúng ta đưa ra kết luận rõ ràng giải thích lý do tại sao các checkpoint Cosmos 16B+ thực tế thường không thể chạy trên phần cứng Colab tiêu chuẩn.

Lập bản đồ gói cosmos-framework

Chúng ta clone và kiểm tra kho lưu trữ cosmos-framework thực tế để hiểu cấu trúc gói, lược đồ đầu vào và quy trình làm việc CLI trực tiếp từ mã nguồn. Chúng ta cũng in ra các mẫu lệnh suy luận chính thức cho việc khởi chạy trên một GPU và nhiều GPU, bao gồm các chế độ như chuyển đổi văn bản thành video (text-to-video), hình ảnh thành video (image-to-video), động lực học tiến (forward dynamics), động lực học nghịch đảo (inverse dynamics) và chính sách (policy). Sau đó, chúng ta giới thiệu ý tưởng Mixture-of-Transformers đa phương thức, nơi các token văn bản, hình ảnh và hành động chia sẻ cơ chế chú ý trong khi vẫn sử dụng các khối feed-forward chuyên gia dành riêng cho từng phương thức.

Xây dựng mô hình Omnimodal MoT

Chúng ta triển khai mô hình Mixture-of-Transformers đa phương thức thu nhỏ từ đầu bằng PyTorch. Chúng ta định nghĩa RMSNorm, rotary embeddings, cơ chế tự chú ý nhân quả (causal self-attention) được chia sẻ, các chuyên gia SwiGLU dành riêng cho từng phương thức và toàn bộ kiến trúc OmniMoT. Sau đó, chúng ta khởi tạo mô hình trên thiết bị khả dụng và báo cáo số lượng tham số, số lớp, cấu trúc chuyên gia và thiết bị runtime.

Huấn luyện trên dữ liệu tổng hợp

Chúng ta tạo một tập dữ liệu thế giới vật lý tổng hợp, nơi các luồng văn bản, hình ảnh và hành động phụ thuộc vào các mã cảnh ẩn. Chúng ta huấn luyện mô hình thế giới thu nhỏ để dự đoán các token văn bản tiếp theo, các trạng thái tiềm ẩn hình ảnh tương lai và các vector hành động tương lai bằng cách sử dụng mục tiêu kết hợp cross-entropy và MSE. Chúng ta theo dõi hàm mất mát huấn luyện qua nhiều bước và tùy chọn vẽ biểu đồ để cho thấy cách mô hình học các động lực học đa phương thức.

Autoregressive World-Model Rollout

Chúng ta kiểm tra mô hình đã huấn luyện bằng cách sử dụng autoregressive rollout, phản ánh phương pháp động lực học tiến được sử dụng trong các mô hình thực tế. Chúng ta cung cấp cho mô hình một tiền tố hình ảnh ngắn và để nó dự đoán các trạng thái tiềm ẩn tương lai từng bước một. Sau đó, chúng ta so sánh quỹ đạo được tưởng tượng với vật lý tổng hợp thực tế và báo cáo MSE để đánh giá mức độ mô hình nắm bắt các động lực học tương lai.

Suy luận Cosmos 3 thực tế

Chúng ta tạo các thông số kỹ thuật đầu vào Cosmos 3 đúng lược đồ cho các tác vụ văn bản thành hình ảnh, văn bản thành video, video có âm thanh, hình ảnh thành video, động lực học tiến và các tác vụ robot kiểu chính sách. Chúng ta cũng in ra các lệnh chính xác cần thiết để khởi chạy suy luận Cosmos 3 thực tế trên phần cứng lớp H100 phù hợp với các thiết lập song song thích hợp. Chúng ta kết thúc bằng việc tóm tắt những gì đã chạy trong Colab và cách quy trình tương tự có thể mở rộng sang các mô hình NVIDIA Cosmos thực tế khi có đủ phần cứng và thiết lập cần thiết.

Kết luận

Tóm lại, chúng ta đã kết nối các ràng buộc thực tế của việc triển khai Cosmos 3 quy mô lớn với một bản triển khai mang tính giáo dục có thể chạy được, giúp chúng ta hiểu về kiến trúc thay vì chỉ đọc về nó. Chúng ta đã huấn luyện một transformer đa phương thức nhỏ, kiểm tra cách các token văn bản, hình ảnh và hành động tương tác thông qua cơ chế chú ý được chia sẻ, và chạy một rollout kiểu động lực học tiến để minh họa khái niệm mô hình hóa thế giới ở quy mô có thể quản lý được. Chúng ta cũng đã tạo các thông số kỹ thuật đầu vào đúng lược đồ và các lệnh khởi chạy chính xác cho suy luận Cosmos 3 thực tế, để quy trình tương tự có thể mở rộng khi chúng ta có quyền truy cập vào phần cứng lớp Ampere hoặc Hopper cần thiết, bộ nhớ GPU đủ lớn, hỗ trợ CUDA và dung lượng ổ đĩa.

Xem toàn bộ mã nguồn với Notebook tại đây. Ngoài ra, hãy thoải mái theo dõi chúng tôi trên Twitter và đừng quên tham gia SubReddit ML 150k+ của chúng tôi và đăng ký nhận bản tin. Khoan đã! Bạn có dùng Telegram không? Bây giờ bạn cũng có thể tham gia cùng chúng tôi trên Telegram.

Cần hợp tác với chúng tôi để quảng bá GitHub Repo, trang Hugging Face, sản phẩm mới hoặc hội thảo trực tuyến của bạn? Hãy kết nối với chúng tôi.

Sana Hassan, thực tập sinh tư vấn tại Marktechpost và là sinh viên bằng kép tại IIT Madras, có niềm đam mê với việc ứng dụng công nghệ và AI để giải quyết các thách thức trong thế giới thực. Với sự quan tâm sâu sắc đến việc giải quyết các vấn đề thực tiễn, cô mang đến một góc nhìn mới mẻ cho sự giao thoa giữa AI và các giải pháp trong đời sống.