Miles v0.1: Hệ thống huấn luyện hậu kỳ cấp sản xuất cho các mô hình AI tiên tiến

Chúng tôi giới thiệu Miles v0.1, một hệ thống full-stack sẵn sàng cho sản xuất (production-ready) dành cho hậu huấn luyện (post-training) các mô hình tiên tiến, là phiên bản kế nhiệm của Miles [1]. Dựa trên thiết kế tinh gọn của slime [16], Miles tối ưu hóa mọi giai đoạn trong vòng lặp huấn luyện RL theo một nguyên tắc đơn giản: xác thực, tinh gọn và có khả năng tùy biến ở mọi nơi. Với các mục tiêu hàng đầu là độ chính xác, hiệu suất, độ tin cậy và khả năng mở rộng, Miles hướng tới việc giúp các nhà nghiên cứu và lập trình viên dễ dàng tiếp cận RL ở quy mô tiên tiến. Trong bài viết này, chúng ta sẽ cùng tìm hiểu Miles từ đầu đến cuối.

Vòng lặp RL của Miles

Một tác vụ huấn luyện RL trong Miles là một vòng lặp qua các giai đoạn sau:

Trong phần tiếp theo, chúng ta sẽ đi qua từng thành phần trong vòng lặp để làm nổi bật cách chúng tôi giúp Miles đạt được độ chính xác, hiệu suất, độ tin cậy và khả năng mở rộng.

Hình 1. Vòng lặp RL hoàn toàn bất đồng bộ của Miles.

Rollout (Triển khai tạo dữ liệu)

Mọi quá trình rollout trong Miles đều được tạo bởi SGLang. Dựa trên hiệu suất suy luận gốc của SGLang, Miles mở khóa quy trình huấn luyện tác nhân (agentic) toàn diện: các phiên đa lượt, thực thi công cụ, môi trường sandbox và ghi lại quỹ đạo (trajectory) trung thực với token.

Rollout tác nhân nhanh bằng SGLang

Miles cung cấp khả năng rollout tác nhân nhanh bằng cách tích hợp trực tiếp SGLang [15], vốn được tối ưu hóa cho việc tạo văn bản dài, đa lượt. Các quỹ đạo tác nhân có độ dài rất khác nhau và mỗi lượt mới đều tái sử dụng phần lớn ngữ cảnh trước đó. Theo mặc định, Miles sử dụng bộ định tuyến SGLang, giúp giữ tất cả các lượt của một phiên trên cùng một engine SGLang và DP rank (nếu DP attention được bật) để tái sử dụng tiền tố (prefix) đã lưu trong bộ nhớ đệm, đồng thời gán các phiên mới cho rank ít tải nhất để tránh việc một vài quỹ đạo dài làm quá tải một phần hệ thống. Bộ định tuyến SGLang cũng dự trữ trước dung lượng KV-cache cho các phiên dài. Những tính năng này đảm bảo sự đồng thời của rollout ổn định và cân bằng, giữ cho tỷ lệ trúng cache cao trong quá trình huấn luyện tác nhân.

RL hoàn toàn bất đồng bộ

Đối với các tác vụ có ngữ cảnh dài, sử dụng công cụ và tác nhân, thời gian rollout chủ yếu bị quyết định bởi một vài thành phần chậm chạp (stragglers). Lịch trình đồng bộ làm trầm trọng thêm vấn đề này vì trainer phải ngồi chờ cho đến khi quỹ đạo chậm nhất trong batch trả về, và các engine rollout phải đợi cho đến khi bộ tối ưu hóa (optimizer) hoàn tất cập nhật mô hình. RL hoàn toàn bất đồng bộ của Miles loại bỏ sự chặn lẫn nhau này bằng cách cho phép các engine rollout tạo dữ liệu liên tục: quá trình tạo rollout diễn ra không ngừng trong khi trainer tiêu thụ các nhóm đã hoàn thành và cập nhật mô hình. Không bên nào bị chặn bởi bên kia.

Hình 2. Cách RL hoàn toàn bất đồng bộ xử lý các đuôi dài (long tails). Các màu khác nhau biểu thị các phiên bản trọng số khác nhau; các đoạn màu xanh lá cây đại diện cho thời gian gọi công cụ.

Lịch trình hoạt động ở mức độ chi tiết của mẫu: mỗi quỹ đạo hoàn thành sẽ giải phóng ngay lập tức một vị trí, giữ cho sự đồng thời của việc tạo dữ liệu ổn định bất chấp sự khác biệt lớn về độ dài quỹ đạo. Các nhóm hoàn thành đi vào một bộ đệm dữ liệu có giới hạn, giúp tách biệt thông lượng rollout khỏi nhịp độ huấn luyện. Bộ đệm này cũng tạo thành một ranh giới chính sách có thể tùy chỉnh, nơi người dùng quyết định mẫu nào cần giữ lại, thử lại, loại bỏ hoặc từ chối vì đã cũ mà không cần sửa đổi lịch trình hoặc thực thi.

Để thực hiện đánh giá bất đồng bộ, Miles cung cấp ba chế độ đánh giá, được phân biệt bởi nguồn gốc của trọng số mô hình.

Cả ba chế độ đều liên kết kết quả với checkpoint và bước huấn luyện đã tạo ra chúng. Nếu quá trình đánh giá kết thúc sau đó vài bước, Miles sẽ báo cáo độ trễ đó thay vì gán sai kết quả; các lỗi sẽ bỏ qua điểm đánh giá thay vì chấm dứt quá trình huấn luyện.

Hình 3. Các chế độ đánh giá: việc chia sẻ các engine rollout sẽ kiểm soát quá trình tạo dữ liệu, trong khi một bản chụp checkpoint sẽ chuyển việc đánh giá sang một đội ngũ chuyên biệt hoặc dịch vụ bên ngoài mà không cần dừng huấn luyện.

Môi trường tác nhân

Trong RL tác nhân, phần lớn công việc diễn ra trong một môi trường cô lập. Ví dụ, huấn luyện một tác nhân lập trình nghĩa là cung cấp cho mô hình sandbox riêng cho từng tác vụ: nó chạy lệnh, chỉnh sửa tệp, đọc kết quả trả về và một bộ kiểm thử ở cuối sẽ quyết định xem tác vụ đã được giải quyết hay chưa. Môi trường lưu giữ trạng thái đó, thực thi các hành động và chạy trình xác thực để tạo ra phần thưởng. Miles chạy nhiều tập (episode) như vậy cùng lúc, ghi lại từng quỹ đạo dưới dạng mà trainer có thể học được và truyền kết quả của trình xác thực dưới dạng phần thưởng.

Một môi trường kết nối với Miles thông qua điểm cắm (plug-in point). Miles cung cấp nhiều điểm cắm ở các cấp độ khác nhau của ngăn xếp rollout, vì vậy một môi trường có thể đảm nhận bao nhiêu phần của rollout tùy thích. Chúng tôi cung cấp các tích hợp cho các hệ sinh thái hiện có, bao gồm Harbor, HUD, NeMo Gym, OpenEnv và Prime Intellect Verifiers. Bạn thậm chí có thể mang theo môi trường của riêng mình và cắm vào thông qua các điểm tương tự.

Các sandbox chạy trên bất kỳ backend nào bạn chọn. Chúng tôi hỗ trợ một số backend bao gồm AgentENV, Daytona, E2B và Modal. Mỗi tập nhận được một sandbox mới, được xây dựng từ image của tác vụ đó, vì vậy không có gì được chia sẻ giữa các tập và không có gì tồn tại sau khi chạy xong.

Miles v0.1 cung cấp các công thức (recipes) end-to-end được duy trì cho các tác vụ lập trình tác nhân và dòng lệnh. Chúng đã được xác thực trên phần cứng thực tế và có thể khởi chạy ngay lập tức.

Token-In-Token-Out (TITO)

Trong RL tác nhân đa lượt, đầu ra của mô hình đi qua quá trình phân tích cú pháp tin nhắn, thực thi công cụ và kết xuất chat-template trước khi bước vào lượt tiếp theo. Quá trình này có thể thay đổi quá trình token hóa, cắt tỉa suy luận lịch sử hoặc tái tuần tự hóa các lệnh gọi công cụ, khiến trainer nhìn thấy một ngữ cảnh token khác với ngữ cảnh thực tế được sử dụng trong quá trình rollout.

Máy chủ phiên TITO trong Miles [6] bảo toàn chính xác các ID token do mô hình tạo ra. Ở mỗi lượt mới, nó chỉ token hóa các tin nhắn mới được thêm vào và hợp nhất chúng với tiền tố hiện có. Điều này cho phép toàn bộ quỹ đạo được tập hợp thành một mẫu huấn luyện liên tục, bảo toàn xác suất log của rollout gốc trong khi thực hiện loss-masking (che giấu mất mát) các token không do mô hình tạo ra. Đối với mỗi dòng mô hình, TITO được xác thực thông qua các bài kiểm tra round-trip trên CPU và các phiên GPU SGLang thực tế, bảo vệ tính chính xác ở cấp độ token cần thiết cho R3, OPD và căn chỉnh zero-KL.

Dựa trên TITO, Miles đang phát triển công thức huấn luyện cho các bộ khung tác nhân hộp đen (black-box) như Claude Code và Codex. Các bộ khung này tạo ra các tác nhân phụ và nén ngữ cảnh tại thời điểm chạy, vì vậy số lượng quỹ đạo mỗi tác vụ là động và không biết trước. Do đó, Miles ghi lại toàn bộ cây quỹ đạo thông qua máy chủ phiên, và ở phía huấn luyện, áp dụng chuẩn hóa mất mát cần thiết cho các kích thước batch thay đổi, giữ cho quy mô gradient nhất quán.

Hình 4. Token-in, token-out. Máy chủ phiên giữ lại chính xác các ID token mà các engine đã tạo ra, vì vậy trainer nhìn thấy các token mà mô hình thực sự tạo ra ngay cả khi bộ khung là một hộp đen mờ đục.

Efficient Rollout Routing Replay (R3)

MoE RL rất nhạy cảm với những khác biệt số học nhỏ giữa rollout và huấn luyện: một lần lật routing top-k duy nhất cũng làm thay đổi cả tính toán token và trọng số chuyên gia (expert weights) nhận gradient. Để giải quyết vấn đề này, Rollout Routing Replay (R3) của Miles ghi lại kết quả định tuyến chuyên gia của SGLang trong quá trình rollout và phát lại chúng trong quá trình huấn luyện. Việc phát lại này được xử lý hiệu quả trong SGLang, chỉ thêm chi phí tối thiểu so với định tuyến thông thường.

Huấn luyện

Trainer là xương sống của RL. Trong Miles, chúng tôi cung cấp nhiều tối ưu hóa trainer giúp RL ổn định, nhanh chóng và tiết kiệm tài nguyên.

Huấn luyện độ chính xác thấp (Low-precision Training)

Miles hỗ trợ rollout ở định dạng NVFP4, MXFP4, MXFP8 và FP8; các công thức huấn luyện end-to-end cho NVFP4, MXFP8 và FP8; và huấn luyện nhận biết lượng tử hóa (QAT) INT4 [3] cho hầu hết các mô hình. Các bài viết trước đã đề cập sâu về các công thức FP8 [2] và INT4 QAT [3].

Việc nắm bắt thông lượng độ chính xác thấp của Blackwell trong RL đòi hỏi nhiều hơn là chỉ hoán đổi các kiểu dữ liệu GEMM: quá trình lượng tử hóa ở phía rollout và huấn luyện phải khớp nhau, nếu không sự sai lệch sẽ tích tụ qua các lần cập nhật trọng số dẫn đến phân kỳ chính sách. Để đạt được điều này, chúng tôi xây dựng các công thức MXFP8 và NVFP4 gốc của Blackwell [9] như các hợp đồng độ chính xác end-to-end trên toàn bộ ngăn xếp: MXFP8 chạy rollout, forward và cả hai gradient GEMM với khả năng mở rộng theo khối phần cứng; NVFP4 lượng tử hóa trọng số chuyên gia MoE theo từng token với khả năng mở rộng kích hoạt trực tuyến để tránh các hiện tượng lượng tử hóa phụ thuộc vào batch; một hợp đồng lượng tử hóa bit-exact giúp các nhân huấn luyện và rollout nhìn thấy cùng các giá trị đã lượng tử hóa, với các cờ chi tiết để giữ các lớp nhạy cảm ở định dạng BF16. Các công thức được xác thực với độ phân kỳ KL thấp giữa SGLang và Megatron-LM. Tất cả các cấu hình độ chính xác thấp đều giữ cho các đường cong phần thưởng bám sát đường cơ sở BF16 trong khi giảm thời gian rollout.

Hiệu quả bộ nhớ & Disk Offload

Huấn luyện một mô hình 744 tỷ tham số một cách bất đồng bộ trên 16 node đặt ra vấn đề bộ nhớ đáng kể. Miles đi kèm với tối ưu hóa bộ nhớ tinh vi mà các lần chạy quy mô lớn phụ thuộc vào, bắt đầu từ bộ tối ưu hóa: các trạng thái bộ tối ưu hóa có thể được offload sang CPU hoặc NVMe cục bộ của node và truyền ngược lại theo từng bucket trong mỗi bước tối ưu hóa, vì vậy chúng không bao giờ cần phải nằm trên GPU cùng một lúc. Trong ví dụ mà chúng tôi trình bày ở cuối bài viết này, việc truyền trạng thái bộ tối ưu hóa qua NVMe này là thứ cho phép bộ tối ưu hóa GLM-5.2 khớp cùng với engine huấn luyện trên 32 GPU GB300. Ngoài việc offload, Miles còn cho phép một loạt các giảm thiểu dấu chân bộ nhớ GPU và CPU khác, tiết kiệm hơn 30 GB bộ nhớ HBM mỗi GPU trong ví dụ chạy GLM-5.2 và hàng trăm GB bộ nhớ CPU mỗi node khi trainer được đặt cùng vị trí với các engine rollout.

Hai Backend huấn luyện

Miles hỗ trợ hai backend huấn luyện đằng sau một giao diện: NVIDIA Megatron-LM và PyTorch FSDP. Một cờ khởi chạy duy nhất sẽ chọn backend nào sở hữu mô hình trên GPU; một backend triển khai năm phương thức và không có gì thay đổi bên trên dòng đó.

Cập nhật trọng số

Trong mỗi bước huấn luyện, sau khi bộ tối ưu hóa cập nhật trọng số mô hình, các trọng số mới phải được đồng bộ hóa trên tất cả các engine rollout. Khi huấn luyện và rollout chạy trên các GPU khác nhau, việc đồng bộ hóa trọng số có thể trở thành nút thắt cổ chai lớn của đường ống. Miles cung cấp hai đường dẫn tối ưu hóa cho các cài đặt triển khai đa dạng. Với truyền trọng số P2P [5], các rank huấn luyện phân đoạn lại từng bucket trọng số cho bố cục SGLang mục tiêu và chỉ ghi các phân đoạn cần thiết trực tiếp vào bộ nhớ rank rollout qua RDMA. Điều này tránh việc phát sóng toàn bộ mô hình đến mọi rank và sử dụng nhiều rank huấn luyện làm bên gửi song song. Lợi ích tăng lên cùng với kích thước mô hình và tính song song của chuyên gia: đối với Kimi-K2 1T, P2P đã giảm thời gian cập nhật trọng số từ 53,3 giây xuống 7,2 giây.