NVIDIA ra mắt Molt: Khung làm việc học tăng cường cho tác tử AI dựa trên PyTorch

Nghiên cứu về agentic reinforcement learning (học tăng cường cho tác nhân) là quá trình sửa đổi thuật toán liên tục. Các bộ ước lượng mới, các giai đoạn pipeline mới, các lược đồ rollout mới. Trong các framework phổ biến, mỗi thay đổi đều phải len lỏi qua các lớp của trainer, backend phân tán và các thành phần kết nối rollout. Chi phí đó đè nặng lên nhà nghiên cứu sau mỗi lần lặp.

Molt, từ nhóm NeMo của NVIDIA, nhắm trực tiếp vào chi phí đó. Đây là một framework agentic RL thuần PyTorch với mục tiêu thiết kế khác thường. Cơ sở mã nguồn phải đủ gọn để một nhà nghiên cứu có thể nắm bắt trong đầu, và đủ để một trợ lý lập trình AI có thể đọc và hiểu toàn bộ. Dung lượng được công bố là khoảng 8,6 nghìn dòng mã RL, được đo bằng cách truy vết biểu đồ import từ điểm bắt đầu RL của mỗi framework. Phương pháp tương tự cho thấy verl có khoảng 62 nghìn dòng, slime là 25 nghìn và OpenRLHF là 7,2 nghìn.

Nó có thể triển khai được không?

Có. Molt được phát hành theo giấy phép Apache 2.0 cùng với các mã khởi chạy, tập lệnh Slurm và một container được xây dựng sẵn. Tuy nhiên, bài báo nghiên cứu định vị nó là cơ sở hạ tầng nghiên cứu chứ không phải dịch vụ huấn luyện sản xuất, và phần cứng mới là rào cản thực sự. Các công thức được cung cấp giả định cấu hình 2 node với 8 GPU H100 mỗi node, chia 8 cho huấn luyện và 8 cho rollout.

Điều này đưa Molt vào tầm với của các phòng thí nghiệm tiên phong hoặc cận tiên phong, các startup AI được tài trợ tốt đang thực hiện hậu huấn luyện (post-training), các nhóm nghiên cứu AI doanh nghiệp trong lĩnh vực tài chính, y tế và robot – những nơi huấn luyện tác nhân trong các môi trường độc quyền, và các phòng thí nghiệm học thuật có quyền truy cập vào hệ thống đa node H100/H200. Các ứng dụng bao gồm tác nhân sử dụng công cụ đa lượt (multi-turn tool-use), tác nhân thực thi mã, môi trường ngôn ngữ-hình ảnh (công thức geo3k đi kèm), vòng lặp phần thưởng LLM-as-judge và chưng cất on-policy sang một mô hình học sinh nhỏ hơn.

Ba thành phần, một vòng lặp

Molt kết hợp Ray để phân bổ và quản lý hàng đợi bất đồng bộ, vLLM cho rollout, và NVIDIA AutoModel với FSDP2 cho huấn luyện. Không có thành phần nào trong ba thành phần này bị fork, vì vậy các cải tiến từ upstream sẽ được cập nhật thông qua việc ghim phiên bản container thay vì phải rebase lại.

Runtime là một pool tác nhân, một tập hợp các engine vLLM nằm sau bộ định tuyến yêu cầu (request router) và một actor chính sách có thể huấn luyện duy nhất. Một pool streaming giữ cho các nhóm prompt luôn hoạt động để các engine không bao giờ bị trống trong khi actor đang huấn luyện. Rollout một phần sẽ tạm dừng các engine, phát các shard của actor qua NCCL trực tiếp đến từng engine và tiếp tục các yêu cầu đang giữ thay vì loại bỏ chúng.

Tác nhân là một chương trình thông thường

Một lần chạy RL sẽ đặt tên cho một module Python xuất ra một AgentRunner. Mọi thứ khác đều là mã thông thường, bao gồm cả phần thưởng.

Hai hình thức được hỗ trợ. Với Env, framework sở hữu vòng lặp LLM trong một step tương thích với Gymnasium. Với ChatAgent, người dùng sở hữu vòng lặp thông qua SDK tiêu chuẩn của OpenAI hoặc Anthropic. Molt khởi chạy một máy chủ loopback hỗ trợ cả hai giao thức truyền tin, và mọi yêu cầu đều được giải mã phía máy chủ thành một chuỗi tích lũy khớp từng token. Khi một tác nhân có tầm nhìn dài hạn nén ngữ cảnh và viết lại tiền tố, máy chủ sẽ đóng phân đoạn hiện tại và tự động mở một phân đoạn mới.

Không bao giờ huấn luyện trên một token mà bạn không tạo ra

Ba bất biến về tính đúng đắn định hình nên thiết kế này. Định danh token: các id token được lấy mẫu xác định quỹ đạo, không phải bản ghi đã được token hóa lại. Ngữ nghĩa phiên bản chính sách: các token có thể huấn luyện giữ lại log-probability của chính sách hành vi, và việc sử dụng bất đồng bộ được hiệu chỉnh trên mỗi token đằng sau một cổng cấp chuỗi. Tính nhất quán tiến (forward consistency): rollout và actor phải đồng nhất về ngữ nghĩa mô hình.

Đối với các chính sách mixture-of-experts, bất biến cuối cùng là quan trọng nhất. Các bộ định tuyến rollout và huấn luyện chọn chuyên gia một cách độc lập, và những khác biệt nhỏ về số học có thể làm thay đổi lựa chọn top-k. Molt áp dụng phát lại định tuyến rollout (rollout routing replay), nơi vLLM trả về các id chuyên gia trên mỗi token và quá trình forward huấn luyện sẽ phát lại chúng.

Giải thích tương tác

Những điểm chính cần lưu ý

Xem bài báo và Repo tại đây. Ngoài ra, hãy 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, Bản phát hành sản phẩm 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.

Asif Razzaq là CEO của Marktechpost Media Inc.. Là một doanh nhân và kỹ sư có tầm nhìn, Asif cam kết khai thác tiềm năng của Trí tuệ nhân tạo vì lợi ích xã hội. Nỗ lực gần đây nhất của ông là ra mắt Nền tảng truyền thông Trí tuệ nhân tạo, Marktechpost, nổi bật với việc đưa tin chuyên sâu về học máy và học sâu, vừa đảm bảo tính kỹ thuật vừa dễ hiểu đối với đông đảo khán giả. Nền tảng này tự hào với hơn 2 triệu lượt xem hàng tháng, minh chứng cho sự phổ biến của nó đối với độc giả.