Mooncake: Giải pháp tối ưu hóa truyền tải dữ liệu Rollout cho hệ thống học tăng cường Miles

Dữ liệu Rollout trong các hệ thống RL phân tách

Học tăng cường (Reinforcement learning) cho các mô hình ngôn ngữ lớn kết hợp hai khối lượng công việc rất khác biệt: tạo rollout và huấn luyện mô hình.

Trong quá trình rollout, các worker suy luận (inference workers) chạy chính sách hiện tại trên một tập hợp các prompt và tạo ra phản hồi. Đồng thời, chúng tạo ra thông tin cần thiết cho thuật toán học, bao gồm các token được tạo, mask, log probability, phần thưởng (rewards), độ dài chuỗi, định danh mẫu và các siêu dữ liệu khác. Tổng hợp lại, các đầu ra này tạo thành dữ liệu rollout sẽ được tiêu thụ bởi bước huấn luyện tiếp theo.

Ở quy mô nhỏ, việc tạo và huấn luyện có thể chia sẻ cùng một môi trường thực thi. Tuy nhiên, ở quy mô lớn hơn, các hệ thống RL hiện đại ngày càng áp dụng kiến trúc phân tách (disaggregated architecture), trong đó việc tạo rollout và huấn luyện được triển khai dưới dạng các nhóm worker riêng biệt, thường trên các tiến trình, GPU hoặc máy chủ khác nhau.

Hai giai đoạn này có đặc điểm về tài nguyên và thực thi cơ bản khác nhau. Rollout là khối lượng công việc nặng về suy luận, có thông lượng phụ thuộc vào hiệu suất giải mã, batching và lập lịch yêu cầu, trong khi huấn luyện dựa trên các phép tính tensor lớn, đồng bộ hóa cao. Việc tách biệt chúng cho phép mỗi bên được mở rộng và lập lịch độc lập thay vì ép buộc cả hai khối lượng công việc vào cùng một mô hình thực thi.

Sự tách biệt này cũng cho phép tính đồng thời ở cấp độ pipeline. Miles có thể huấn luyện trên rollout N trong khi các worker rollout đã bắt đầu tạo rollout N+1. Do đó, các hệ thống RL bất đồng bộ có thể cho phép các worker rollout và huấn luyện tiến hành với tốc độ khác nhau thay vì bắt một giai đoạn phải đợi giai đoạn kia sau mỗi thao tác. Lợi ích là tận dụng tài nguyên tốt hơn và linh hoạt hơn trong việc cung cấp năng lực suy luận và huấn luyện. Nhưng việc phân tách cũng tạo ra một ranh giới hệ thống mới: dữ liệu do các worker rollout tạo ra giờ đây phải di chuyển đến một nhóm worker khác trước khi quá trình huấn luyện có thể tiêu thụ nó.

Việc bàn giao đó nằm trực tiếp giữa quá trình tạo và cập nhật chính sách tiếp theo. Việc chuyển dữ liệu chậm có thể khiến các trainer phải chờ đợi và giữ bộ nhớ trên các worker rollout lâu hơn mức cần thiết. Đối với một hệ thống RL phân tán như Miles, việc di chuyển hiệu quả dữ liệu rollout có cấu trúc này từ phía suy luận sang phía huấn luyện trở thành một phần quan trọng của pipeline RL đầu cuối.

Điều gì làm cho việc truyền dữ liệu RL Rollout trở nên thách thức

Dữ liệu rollout RL khác biệt đáng kể so với các tensor lớn, đều đặn thường được di chuyển trong huấn luyện phân tán.

Một batch rollout thường là một đối tượng có cấu trúc không đồng nhất thay vì một tensor liên tục duy nhất. Tùy thuộc vào framework và thuật toán, nó có thể chứa các token được tạo, loss mask, log probability, phần thưởng, độ dài chuỗi, định danh mẫu, thông tin định tuyến, siêu dữ liệu và các trường phụ trợ khác. Các giá trị này có thể được biểu diễn dưới dạng tensor, mảng NumPy, danh sách vô hướng Python, mảng có độ dài thay đổi theo từng mẫu, byte hoặc các đối tượng Python tùy ý.

Một số đặc tính khiến dữ liệu này đặc biệt khó di chuyển hiệu quả.

Thách thức 1: Các kiểu dữ liệu không đồng nhất và ngữ nghĩa phức tạp

Các trường rollout khác nhau có biểu diễn và ngữ nghĩa cơ bản khác nhau. Các tensor dày đặc và mảng số có thể được truyền hiệu quả dưới dạng bộ đệm có kiểu (typed buffers), trong khi các chuỗi không đều (ragged sequences) cần thông tin về ranh giới hàng, danh sách vô hướng cần giữ nguyên giá trị gốc, và siêu dữ liệu hoặc đối tượng Python có thể yêu cầu mã hóa tổng quát hơn. Đồng thời, trainer phải tái tạo chính xác cấu trúc mà framework RL mong đợi, bao gồm dtype, shape, thứ tự hàng, trạng thái null và siêu dữ liệu. Một bộ tuần tự hóa (serializer) chung có thể xử lý các đối tượng này về mặt chức năng, nhưng thường phải trả giá bằng việc chuyển đổi, sao chép và tái tạo thêm. Do đó, việc truyền rollout hiệu quả cần hiểu cả biểu diễn vật lý và cấu trúc logic của từng trường.

Thách thức 2: Bố cục bộ nhớ bị phân mảnh cao

Dữ liệu rollout có thể chứa một số lượng rất lớn các phân bổ bộ nhớ nhỏ. Trong khối lượng công việc Miles đã ghi lại, các trường chính như tokens, loss_masks và rollout_log_probs được biểu diễn dưới dạng list[np.ndarray], với mỗi mảng NumPy được cấp phát cho mỗi mẫu. Khi kích thước batch tăng lên, việc truyền các mảnh này riêng lẻ dẫn đến việc đăng ký bộ nhớ và các thao tác Store lặp đi lặp lại, trong khi việc tuần tự hóa toàn bộ đối tượng Python đòi hỏi phải duyệt, sao chép và xây dựng lại một đồ thị đối tượng lớn. Thách thức là biến dữ liệu logic bị phân mảnh thành các lần truyền khối hiệu quả mà không làm mất cấu trúc ban đầu của nó.

Tổng hợp lại, những thách thức này làm cho việc di chuyển dữ liệu rollout không chỉ là vấn đề băng thông: hệ thống phải di chuyển hiệu quả dữ liệu phân mảnh, không đồng nhất trong khi vẫn bảo toàn cấu trúc của nó.

Một đường dẫn dữ liệu hiệu quả cần đáp ứng đồng thời một số yêu cầu:

Cung cấp năng lượng cho việc truyền dữ liệu Miles Rollout với Mooncake

Miles là một framework học tăng cường hiệu năng cao cho việc hậu huấn luyện mô hình quy mô lớn. Nó kết hợp SGLang để tạo rollout thông lượng cao với Megatron-LM để huấn luyện có khả năng mở rộng, đồng thời cung cấp backend PyTorch FSDP2 cho các khối lượng công việc ưu tiên huấn luyện trực tiếp các triển khai mô hình Hugging Face. Miles hỗ trợ RL hoàn toàn bất đồng bộ, nơi các worker rollout và huấn luyện được tách rời và có thể tiến hành độc lập, cùng với các tính năng như cập nhật trọng số nhanh trong vòng lặp, rollout tác tử (agentic rollout), huấn luyện độ chính xác thấp và khả năng chịu lỗi cho các khối lượng công việc RL quy mô lớn trong sản xuất.

Thiết kế phân tách và bất đồng bộ này làm cho đường dẫn dữ liệu từ rollout sang huấn luyện trở thành một phần quan trọng của pipeline RL. Các batch rollout có cấu trúc và không đồng nhất, thường chứa dữ liệu phân mảnh theo từng mẫu và siêu dữ liệu đặc thù của framework, khiến việc truyền và tái tạo hiệu quả ngày càng trở nên quan trọng khi quy mô khối lượng công việc tăng lên.

Mooncake cung cấp một mặt phẳng dữ liệu (data plane) hiệu năng cao cho các khối lượng công việc AI phân tán. Đối với dữ liệu rollout RL, nó mở rộng mặt phẳng dữ liệu này với khả năng truyền đối tượng có cấu trúc, cho phép các đối tượng rollout không đồng nhất và bị phân mảnh được di chuyển trong khi vẫn bảo toàn cấu trúc và ngữ nghĩa ban đầu của chúng.

Mooncake hiện đã được tích hợp vào Miles như một backend truyền dữ liệu rollout. Trên dữ liệu rollout thu thập từ Miles, sự tích hợp này mang lại độ trễ truyền thấp hơn đáng kể so với đường dẫn Ray hiện có: remote GET nhanh hơn khoảng 10–14 lần, trong khi PUT cải thiện khoảng 1,2–1,6 lần.

Kết quả là quá trình bàn giao từ rollout sang huấn luyện nhanh hơn mà không làm thay đổi mô hình lập trình RL, cung cấp cho Miles một đường dẫn dữ liệu hiệu quả hơn cho các khối lượng công việc rollout lớn, có cấu trúc.

Cách dữ liệu Rollout di chuyển qua Pipeline RL

Miles hỗ trợ cả vòng lặp huấn luyện đồng bộ và bất đồng bộ. Trong bất kỳ chế độ nào, khi các worker rollout và trainer được triển khai trong các tiến trình riêng biệt hoặc trên các máy chủ riêng biệt, mỗi batch rollout hoàn thành phải vượt qua ranh giới đó trước khi có thể được tiêu thụ bởi bước huấn luyện tiếp theo.

Một nguyên tắc thiết kế hữu ích là tách mặt phẳng điều khiển (control plane) khỏi mặt phẳng dữ liệu (data plane). Bộ lập lịch framework quyết định nơi một batch rollout nên đi đến, nhưng nó không nên tự mang theo tải trọng lớn. Thay vào đó, nó truyền một tham chiếu truyền tải nhẹ (lightweight transfer reference) loại bỏ các tải trọng tensor và bố cục chunk của Store, trong khi vẫn cho phép siêu dữ liệu an toàn với JSON khi cần. Tải trọng rollout thực tế đi theo một đường dẫn riêng biệt:

Sự tách biệt này giữ cho các quyết định lập lịch nhẹ nhàng trong khi cho phép tải trọng rollout lớn di chuyển qua một đường dẫn dữ liệu chuyên dụng.

Những gì Miles thực sự truyền tải

Dữ liệu rollout thu thập từ Miles làm cho đường dẫn dữ liệu này trở nên cụ thể. Đối với một cấu hình framework nhất định, hợp đồng trường (field contract) là ổn định, nhưng các trường không chia sẻ một biểu diễn trong bộ nhớ thuận tiện. Chúng được chia thành ba nhóm lớn:

Nhóm đầu tiên mang hầu hết các byte trong dữ liệu thu thập của chúng tôi. Các trường khác nhỏ hơn, nhưng chúng không thể bị loại bỏ hoặc chuẩn hóa: chúng mang định danh mẫu, độ dài, phần thưởng, trạng thái đặc trưng và thông tin kế toán. Kích thước chính xác của chúng phụ thuộc vào khối lượng công việc; phần benchmark đưa ra một ví dụ đã đo lường.

Miles sử dụng cơ chế bàn giao Completed-Dict

Trong Miles, quá trình bàn giao hiện tại bắt đầu sau khi từ điển rollout hoàn chỉnh đã sẵn sàng. Producer gọi put(data, type="dict"), bộ lập lịch mang theo tham chiếu được trả về, và trainer gọi get để tái tạo từ điển gốc.

Hình 1. Miles sử dụng put và get đồng bộ cho một từ điển rollout đã hoàn thành. Tham chiếu đi qua bộ lập lịch, trong khi Mooncake di chuyển tải trọng qua mặt phẳng dữ liệu Store.

Các lệnh gọi truyền tải riêng lẻ là đồng bộ. Producer chỉ trả về tham chiếu sau khi put hoàn tất, và trainer đợi get kết thúc trước khi tiêu thụ đối tượng.

Điều này không ngăn cản bản thân pipeline RL trở nên bất đồng bộ. Miles có thể tạo rollout N+1 trong khi huấn luyện trên rollout N. Nói cách khác, tính đồng thời nằm phía trên thao tác truyền tải: mỗi lần bàn giao riêng lẻ là đồng bộ, trong khi các giai đoạn rollout và huấn luyện khác nhau có thể chồng lấp ở cấp độ pipeline.

Cách Mooncake bảo toàn và truyền dữ liệu Miles Rollout

Mooncake giải quyết hai thách thức ở các lớp khác nhau. Nó giữ cho cấu trúc hiển thị đủ lâu để chọn một biểu diễn hiệu quả cho từng trường: các tensor và mảng vẫn giữ kiểu, các hàng không đều mang siêu dữ liệu ranh giới nhỏ gọn, và các giá trị Python giữ lại những gì GET cần để xây dựng lại chúng. Sau đó, nó biến bộ nhớ bị phân mảnh thành I/O khối. Một kế hoạch sao chép đóng gói các hàng nhỏ đủ điều kiện trực tiếp vào các chunk BufferPool có thể tái sử dụng và đã đăng ký, trong khi các tensor và mảng lớn liên tục có thể sử dụng các đường dẫn Store gốc. Điều này tránh được cả hai thái cực: tuần tự hóa toàn bộ từ điển thành một khối mờ hoặc thực hiện thao tác Store cho mỗi phân bổ nhỏ.

Các thành phần tải trọng kết quả được xuất bản dưới dạng một gói (bundle). Bản kê khai (manifest) của nó ghi lại vị trí của các thành phần đó và cách chúng khớp với nhau, và chỉ trở nên hiển thị sau khi tải trọng đã sẵn sàng. GET tuân theo mô tả đó theo chiều ngược lại để tìm nạp và tái tạo từ điển Miles gốc. Miles vẫn sử dụng giao diện công khai nhỏ, put(data, type="dict") và get(ref, type="dict"); lược đồ, bố cục trường, kế hoạch truyền tải và tái tạo vẫn nằm bên trong Mooncake.

Hình 2. Lược đồ và việc mở rộng lá (leaf expansion) làm lộ kiểu và cấu trúc của từng trường. Mã hóa đặc thù cho từng trường tạo ra các thành phần tải trọng có kiểu và siêu dữ liệu tái tạo; Bundle Store xuất bản bản kê khai của chúng sau cùng. Các lần truyền phân mảnh đủ điều kiện sử dụng staging dựa trên BufferPool, và GET tuân theo cùng cấu trúc đó theo chiều ngược lại.