DeepSeek-V4 Flash triển khai huấn luyện RL trên GPU AMD Instinct MI355X thông qua khung Miles

DeepSeek-V4 RL hiện đã được hỗ trợ trong Miles trên các GPU AMD Instinct™ MI355X với ROCm™! RL yêu cầu quá trình rollout của SGLang và quá trình huấn luyện của Megatron phải thực hiện cùng một chính sách đủ sát sao để xác suất token vẫn được đồng bộ, ngay cả khi Miles liên tục chuyển các trọng số đã cập nhật trở lại công cụ rollout trực tuyến.

DeepSeek-V4 Flash khiến việc này trở nên thách thức thông qua cơ chế attention nén lai, trộn residual mHC và định tuyến MoE. Quá trình triển khai của chúng tôi đã đồng bộ hóa hành vi mô hình giữa SGLang và Megatron, bảo toàn trạng thái lượng tử hóa trong quá trình cập nhật trọng số trực tuyến và thiết lập quy trình làm việc bốn nút (four-node) từ đầu đến cuối. Chúng tôi đã xác thực độ chính xác bằng chênh lệch log-probability giữa huấn luyện và rollout có giới hạn, và trong một lần chạy mở rộng, điểm chuẩn AIME-2024 ngoại tuyến tăng dần cùng với phần thưởng trực tuyến được cải thiện.

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

DeepSeek-V4 Flash trong một hình ảnh

DeepSeek-V4 Flash là mô hình MoE với 284 tỷ tham số, trong đó có 13 tỷ tham số hoạt động trên mỗi token. Cấu hình được sử dụng trong công trình này có 43 lớp decoder, 256 chuyên gia (experts) được định tuyến với lựa chọn top-6, bốn luồng residual mHC và cơ chế attention lai kết hợp cửa sổ trượt 128-token với attention ngữ cảnh dài được nén.

Các lớp C4 chọn 512 mục hàng đầu từ chuỗi KV được nén tỷ lệ 4:1, trong khi các lớp C128 thực hiện attention dày đặc trên chuỗi được nén tỷ lệ 128:1. Cùng với mHC và định tuyến MoE, đây là những đường dẫn cụ thể theo kiến trúc mà SGLang và Megatron phải triển khai một cách nhất quán.

Hình 1. Khối DeepSeek-V4 đơn giản hóa hiển thị quá trình trộn residual mHC, attention nén lai và định tuyến MoE top-6.

Hệ thống Miles

Miles điều phối vòng lặp bất đồng bộ. SGLang tạo ra các phản hồi ứng viên và log-probability của rollout; Megatron chấm điểm các chuỗi tương tự, tính toán cập nhật chính sách và huấn luyện tác nhân (actor); sau đó Miles chuyển các trọng số đã cập nhật trở lại các worker SGLang đang hoạt động.

Đường dẫn FP8 hiện tại sử dụng checkpoint FP8 của Hugging Face cho rollout và checkpoint BF16 torch_dist của Megatron để huấn luyện tác nhân. Do đó, hai công cụ này đại diện cho một chính sách trong các định dạng thực thi khác nhau, khiến việc chuyển đổi, chấm điểm lại và hành vi cập nhật trực tuyến trở thành một phần của ranh giới tính đúng đắn.

Hình 2. Các prompt chảy đến rollout của SGLang; các quỹ đạo và log-probability chảy qua Miles đến tác nhân Megatron; các trọng số đã cập nhật quay trở lại SGLang trước lần rollout tiếp theo.

Thách thức 1: thu hẹp khoảng cách log-probability giữa huấn luyện và rollout

Huấn luyện RL phụ thuộc vào việc SGLang và Megatron gán các xác suất tương tự cho cùng các token được tạo ra. Chúng tôi đã xây dựng một quy trình so sánh đồng nhất token: SGLang tạo một chuỗi một lần, cả hai công cụ đều chấm điểm cùng các token đó, và những khác biệt ở cấp độ token sẽ tiết lộ các sai lệch ở cấp độ mô hình trước khi thực hiện xác thực đa nút tốn kém.

Quá trình so sánh này đã xác định sự khác biệt trong hai đường dẫn cụ thể của DeepSeek-V4: MoE định tuyến bằng hash sớm và trộn residual mHC. Chúng tôi đã căn chỉnh hành vi định tuyến hash của Megatron với SGLang và sửa lỗi mHC post-mix phía Megatron để cả hai công cụ đều bảo toàn cùng một ngữ nghĩa mô hình. Những thay đổi có mục tiêu này đã giúp rollout và huấn luyện đạt được sự đồng thuận về số liệu gần hơn.

Thách thức 2: bảo toàn ngữ nghĩa lượng tử hóa trong quá trình cập nhật trực tuyến

Trong RL, máy chủ rollout nhận các trọng số chính sách đã cập nhật liên tục mà không cần khởi động lại. Đối với các mô hình lượng tử hóa, việc chuyển đổi thành công là chưa đủ: các trọng số đã đóng gói, tensor tỷ lệ và trạng thái runtime phụ thuộc vào lượng tử hóa phải giữ nguyên ý nghĩa dự định của chúng.

Đối với các tensor FP4 và E8M0, AMD đã làm cho đường dẫn cập nhật nhận biết được kiểu dữ liệu, ngăn chặn việc hiểu sai tensor sau cập nhật dẫn đến các thế hệ (generations) không hợp lệ. Đối với FP8, Miles đã xác định vòng đời sau cập nhật; AMD đã thêm giao diện SGLang còn thiếu vào stack ROCm để Miles có thể chạy quá trình xử lý lượng tử hóa cần thiết trước khi rollout tiếp tục.

Bài học chính rất đơn giản: cập nhật trực tuyến phải khôi phục trạng thái lượng tử hóa của mô hình, chứ không chỉ đơn thuần là sao chép các byte của nó.

Thách thức 3: chiến lược song song đa nút ổn định trên ROCm

Việc mở rộng DeepSeek-V4 Flash RL trên nhiều nút AMD Instinct MI355X đã làm nảy sinh hai vấn đề triển khai liên quan: chọn chiến lược song song mô hình phù hợp với mô hình MoE 284 tỷ tham số ở ngữ cảnh 4K và giữ cho giao tiếp tập thể đa nút ổn định. Hai vấn đề này có liên quan với nhau. Tính song song tensor nặng hơn làm giảm bộ nhớ trên mỗi GPU nhưng tăng lưu lượng giao tiếp tập thể, và một số cấu hình đa nút ban đầu đã bị đình trệ bên trong các tập thể RCCL - ví dụ như all-reduce song song tensor hoặc all-to-all chuyên gia không hoàn thành và bị trình theo dõi giao tiếp bắt được, làm dừng quá trình chạy.

Chúng tôi đã hội tụ về một bố cục vừa khả thi về bộ nhớ vừa ổn định: song song tensor 1, song song pipeline 4 và song song chuyên gia 4 trên bốn nút tám GPU, kết hợp với tính toán lại kích hoạt (activation recomputation), giảm tải trạng thái bộ tối ưu hóa (optimizer-state) vào bộ nhớ host và ngân sách token giới hạn trên mỗi GPU. Việc chuyển tính song song từ all-reduce song song tensor sang song song pipeline và chuyên gia, cùng với các cài đặt truyền tải RCCL đã được tinh chỉnh, đã cho phép quá trình chạy diễn ra từ đầu đến cuối trong hơn 100 bước tối ưu hóa mà không bị đình trệ tập thể. Việc thiết lập điểm vận hành ổn định này là điều kiện tiên quyết cho quá trình xác thực dài hơn sau đó.

Xác thực bốn nút trên GPU AMD Instinct MI355X

Chúng tôi đã xác thực đường dẫn FP8 trên bốn nút AMD Instinct MI355X tám GPU: hai nút cho rollout SGLang và hai nút cho huấn luyện tác nhân Megatron. Miles đã điều phối quá trình huấn luyện kiểu GRPO trên khối lượng công việc toán học ngữ cảnh dài (DAPO-Math-17K ở ngữ cảnh 4K), thu thập phần thưởng và cập nhật trọng số trực tuyến lặp đi lặp lại, sử dụng cấu hình song song mô hình gồm song song tensor 1 / song song pipeline 4 / song song chuyên gia 4. Cứ mười bước, chúng tôi lại chạy đánh giá ngoại tuyến trên AIME-2024 với tám mẫu cho mỗi bài toán. Mô hình rollout sử dụng FP8, trong khi tác nhân huấn luyện bằng BF16.

Tính đúng đắn và phần thưởng trực tuyến

Một kiểm tra tính đúng đắn quan trọng là liệu rollout và huấn luyện có gán các xác suất tương tự cho cùng các token được tạo ra hay không. Qua các bước đã ghi lại, chênh lệch log-probability tuyệt đối trung bình là ~0,09. Như Hình 3 cho thấy, sự khác biệt vẫn nằm trong giới hạn qua hơn 100 bước và các lần cập nhật trọng số lặp đi lặp lại, không có sự trôi dạt tăng dần bền vững hoặc tăng đột biến sau khi cập nhật. Đây là một kết quả triển khai đáng khích lệ, không phải là ngưỡng cuối cùng.

Hình 3. Chênh lệch log-probability tuyệt đối giữa huấn luyện và rollout trong 100 bước huấn luyện đầu tiên.

Ngoài sự đồng thuận log-probability có giới hạn, phần thưởng trực tuyến cũng được cải thiện trong quá trình huấn luyện. Trong một lần chạy mở rộng, phần thưởng thô trực tuyến cho thấy xu hướng tăng rõ rệt thay vì đi ngang: giá trị trung bình của nó tăng từ một phần ba đầu tiên của quá trình chạy đến một phần ba cuối cùng (Hình 4). Điều này cho thấy tác nhân đang cải thiện dưới quá trình huấn luyện GRPO liên tục và các lần cập nhật trọng số trực tuyến lặp đi lặp lại, chứ không chỉ duy trì phần thưởng.

Hình 4. Phần thưởng thô trực tuyến qua 100 bước huấn luyện (giá trị mỗi bước, trung bình trượt và khớp tuyến tính). Phần thưởng tăng dần trong suốt quá trình chạy, với độ dốc tuyến tính dương.

Đánh giá ngoại tuyến trên AIME-2024

Tỷ lệ vượt qua trực tuyến được đo lường trên khối lượng công việc huấn luyện và bị ảnh hưởng bởi lấy mẫu động, vì vậy chúng tôi cũng đã chạy một điểm chuẩn ngoại tuyến độc lập, AIME-2024 với tám mẫu cho mỗi bài toán, cứ sau mười bước. Đây là thước đo trung thực về chất lượng mô hình. Trong 100 bước đầu tiên, pass@1 ngoại tuyến của AIME đã cải thiện từ 0,39 lên 0,49 và pass@8 từ 0,53 lên 0,67, trong khi việc cắt ngắn phản hồi ở giới hạn 4.096 token giảm từ 60% xuống 55%. Độ chính xác đơn lẻ và độ bao phủ đa mẫu cùng cải thiện, cho thấy khả năng thực sự đạt được dưới GRPO thay vì chỉ đơn thuần là tinh chỉnh. Các giá trị trên mỗi lần đánh giá có độ nhiễu trên một điểm chuẩn 30 bài toán, vì vậy xu hướng, chứ không phải bất kỳ điểm đơn lẻ nào, mới là tín hiệu.

Hình 5. Pass@1/2/4/8 ngoại tuyến của AIME-2024 qua 100 bước huấn luyện RL đầu tiên (đánh giá mỗi 10 bước, tám mẫu mỗi bài toán). Cả độ chính xác đơn lẻ (pass@1) và độ bao phủ (pass@8) đều có xu hướng tăng.

Những gì chúng tôi đã học được

Đánh giá điểm chuẩn ngoại tuyến là tín hiệu trung thực. Tỷ lệ vượt qua trực tuyến trên khối lượng công việc huấn luyện bị ảnh hưởng bởi lấy mẫu động và lấy mẫu quá mức; đánh giá AIME-2024 độc lập mỗi mười bước đã đưa ra thước đo đáng tin cậy về chất lượng mô hình. Chúng tôi khuyên bạn nên kết hợp mọi lần chạy RL với đánh giá ngoại tuyến định kỳ.

RL đã cải thiện cả độ chính xác và độ bao phủ. Qua 100 bước, pass@1 ngoại tuyến của AIME tăng từ 0,39 lên 0,49 và pass@8 từ 0,53 lên 0,67. Sự gia tăng đồng thời của pass@1 và pass@k cho thấy chính sách đã đạt được khả năng, chứ không chỉ đơn thuần là tinh chỉnh xung quanh các giải pháp mà nó đã có.

Sự đồng thuận giữa các công cụ ổn định trong các lần chạy dài. Chênh lệch log-probability giữa huấn luyện và rollout vẫn nằm trong giới hạn qua hơn 100 bước và các lần cập nhật trọng số lặp đi lặp lại, xác nhận rằng sự căn chỉnh triển khai vẫn giữ vững tốt ngoài cửa sổ xác thực ban đầu.

Việc cắt ngắn phản hồi là giới hạn chính đối với điểm đánh giá tuyệt đối. Khoảng 55-60% phản hồi AIME đạt giới hạn tạo 4.096 token; tăng ngân sách phản hồi đánh giá là đòn bẩy hiệu quả nhất để đạt được độ chính xác tuyệt đối cao hơn.

Đọc xu hướng, không phải các bước đơn lẻ. Phần thưởng trực tuyến mỗi bước và tỷ lệ vượt qua mỗi lần đánh giá có độ nhiễu (phần thưởng trực tuyến dao động 0,36-0,77 theo từng bước); trung bình trượt và đánh giá ngoại tuyến định kỳ là các tín hiệu tiến triển đáng tin cậy.

Hướng đi tiếp theo

Các lệnh khởi chạy