Miles tối ưu hóa học tăng cường trên kiến trúc Blackwell với MXFP8 và NVFP4

TL;DR: Chúng tôi đã triển khai hai công thức RL (Học tăng cường) dựa trên kiến trúc Blackwell trong Miles: MXFP8 đầu-cuối (end-to-end) và NVFP4 theo từng token cho các chuyên gia MoE (Mixture of Experts). Cả hai đều được hỗ trợ bởi khả năng kiểm soát độ chính xác chi tiết trong quá trình chuyển đổi checkpoint, huấn luyện Megatron, triển khai SGLang và cập nhật trọng số trực tiếp. MXFP8 bao phủ các giai đoạn rollout, lan truyền tiến (forward propagation), GEMM trọng số-gradient và GEMM dữ liệu-gradient. NVFP4 sử dụng tính năng điều chỉnh tỷ lệ kích hoạt (activation scaling) theo từng token trực tuyến cho đường dẫn chuyên gia MoE, và cả hai định dạng đều hỗ trợ các chế độ lan truyền ngược (backward) độ chính xác cao hoặc giải lượng tử hóa (dequantized). Trong một thử nghiệm cắt bỏ (ablation) với công thức Qwen3-30B-A3B trên 8x B200, BF16 và tất cả năm cấu hình độ chính xác thấp cho thấy các đường cong phần thưởng thô chồng lấp chặt chẽ, trong khi MXFP8 và NVFP4 giúp giảm thời gian rollout.

Giới thiệu

Trong RL độ chính xác thấp, các giai đoạn rollout, huấn luyện, chuyển đổi checkpoint và cập nhật trọng số trực tiếp phải thống nhất về một hợp đồng độ chính xác (precision contract), nếu không các chính sách của bộ lấy mẫu (sampler) và bộ huấn luyện (trainer) sẽ bị phân kỳ. Miles và hệ sinh thái RL của SGLang đã tích hợp các công thức độ chính xác thấp: bài viết về LMSYS FP8 RL cho thấy việc sử dụng FP8 trong cả huấn luyện và lấy mẫu giúp giảm sự sai lệch so với huấn luyện BF16 với rollout FP8; bài viết về INT4 QAT cho thấy việc lượng tử hóa giả (fake quantization) trong quá trình huấn luyện và rollout W4A16 có thể làm cho INT4 trở nên thực tế. Chúng tôi mở rộng công việc đó sang các định dạng gốc của Blackwell bằng cách triển khai các công thức MXFP8 và NVFP4 trong Miles và đóng góp các thành phần hỗ trợ trên SGLang, TransformerEngine, FlashInfer, Megatron và giao diện cuDNN. Lộ trình công khai được theo dõi tại miles#615.

Các đóng góp chính bao gồm:

Tại sao lại cần các công thức dựa trên kiến trúc Blackwell?

Các phương pháp độ chính xác thấp trước đây không được thiết kế xoay quanh MXFP8 hoặc NVFP4. Đường dẫn Miles hiện tại tuân theo công thức FP8 chia khối (block-scaled) kiểu DeepSeek-V3: trọng số sử dụng chia tỷ lệ khối 128x128, kích hoạt sử dụng chia tỷ lệ ô 1x128, và các tỷ lệ được tính toán trực tuyến cho mỗi ô hoặc khối. Đây là một công thức mạnh mẽ thời Hopper, nhưng trên Blackwell, các tỷ lệ FP32 của nó vẫn được áp dụng bằng phần mềm xung quanh đường dẫn Tensor Core thay vì thông qua phần cứng vi chia tỷ lệ (microscaling) gốc.

INT4 QAT giải quyết một vấn đề khác. Huấn luyện sử dụng lượng tử hóa giả để thích ứng mô hình với trọng số INT4, trong khi rollout sử dụng W4A16. Mặc dù tiết kiệm bộ nhớ, đường dẫn tính toán vẫn thực sự sử dụng các kích hoạt BF16 với trọng số INT4 đã giải lượng tử hóa. Bảng dưới đây chuẩn hóa các thông số kỹ thuật Tensor Core dày đặc của nền tảng HGX của NVIDIA theo thông lượng trên mỗi GPU: B200 và B300 từ các hệ thống HGX 8-GPU, Rubin từ bảng HGX Rubin NVL8.

Đối với các hệ thống RL, hợp đồng độ chính xác bao gồm:

Bối cảnh định dạng

MXFP8

MXFP8 là một định dạng FP8 vi chia tỷ lệ. Tài liệu MXFP8 của TransformerEngine mô tả nó là một công thức chia tỷ lệ theo khối gốc của Blackwell: cứ mỗi 32 giá trị E4M3 liên tiếp sẽ chia sẻ một tỷ lệ E8M0 cục bộ, và khối này là một chiều.

Vì các tỷ lệ E8M0 đại diện cho lũy thừa của hai, tỷ lệ đã giải mã thường được làm tròn lên để giá trị tối đa trong khối không bị cắt xén (clipped).

NVFP4

NVFP4 là định dạng FP4 gốc của Blackwell. Như đã mô tả trong phần giới thiệu NVFP4 của NVIDIA, nó lưu trữ các giá trị FP4 E2M1 với một tỷ lệ FP8 E4M3 cho mỗi khối 16 giá trị. Vì E4M3 có độ phân giải tốt hơn UE8M0, tỷ lệ của nó thường được làm tròn đến giá trị có thể biểu diễn gần nhất. Một công thức NVFP4 tiêu chuẩn cũng thêm một tỷ lệ FP32 cho phạm vi tensor lớn hơn, tạo ra một hệ thống phân cấp hai cấp:

Tỷ lệ FP32 có thể được chọn trên các phạm vi tensor khác nhau. Lựa chọn đó là một quyết định về công thức thay vì là thuộc tính của chính định dạng, và nó trở nên đặc biệt quan trọng đối với RL.

Công thức 1: RL MXFP8 đầu-cuối (End-to-End)

Công thức MXFP8 là phần mở rộng trực tiếp nhất dựa trên kiến trúc Blackwell của công việc FP8 đầu-cuối trước đó. Rollout, lan truyền tiến, GEMM trọng số-gradient và GEMM dữ liệu-gradient đều sử dụng MXFP8, trong khi các tensor được chọn vẫn giữ ở định dạng BF16 thông qua các quy tắc kiểm soát độ chính xác được mô tả bên dưới.

Huấn luyện

TransformerEngine và Megatron triển khai MXFP8 như một đường dẫn huấn luyện Blackwell hạng nhất, được tối ưu hóa hiệu suất, bao gồm các tối ưu hóa GB200 DeepSeek-V3 được mô tả trong deepseek-v3-gb200-optimization.md. Trong tích hợp Miles của chúng tôi, chúng tôi sử dụng đường dẫn này làm nền tảng phía huấn luyện cho RL MXFP8 đầu-cuối.

Một điểm khác biệt so với công thức FP8 của DeepSeek-V3 là cách biểu diễn các kích hoạt ngược (backward activations). DeepSeek-V3 lưu trữ các kích hoạt tiến trong các ô FP8 1x128 và chuyển đổi chúng sang hướng ngược trước khi thực hiện GEMM ngược. Cách tiếp cận đó lưu trữ ít dữ liệu FP8 hơn, nhưng nó giới thiệu quá trình giải lượng tử hóa cộng với tái lượng tử hóa trước GEMM ngược. Tài liệu MXFP8 của TransformerEngine lưu ý rằng các khối 1x32 theo hàng và các khối 32x1 theo cột có sự khác biệt về mặt số học và phải được lượng tử hóa độc lập với dữ liệu độ chính xác đầy đủ. Do đó, TransformerEngine hiện thực hóa cả bản sao lượng tử hóa theo hàng và theo cột trong quá trình lượng tử hóa. Điều này sử dụng nhiều bộ nhớ hơn, nhưng nó tránh được một bước tái lượng tử hóa bổ sung và giảm sai số lượng tử hóa thêm trong đường dẫn ngược.

Đây là một sự đánh đổi hệ thống điển hình cho RL. Chúng tôi sử dụng đường dẫn TransformerEngine để duy trì một hợp đồng MXFP8 đầu-cuối mà không tạo thêm nguồn sai lệch nào khác.

Rollout

Về phía rollout, SGLang sử dụng các nhân (kernels) Blackwell MXFP8 từ FlashInfer và Triton. Chúng tôi đã triển khai và đưa đường dẫn rollout lên các kho lưu trữ FlashInfer và SGLang (flashinfer#2581, sglang#17449, sglang#19537, sglang#21576, và sglang#28459).

Hầu như tất cả các GEMM chính đều có thể được lượng tử hóa sang MXFP8, bao gồm các phép chiếu attention và các chuyên gia MoE. Các ngoại lệ chính là các lớp độ chính xác cao được kiểm soát rõ ràng, chẳng hạn như các phép chiếu BF16 MLA được mô tả bên dưới.

Công thức 2: RL NVFP4 theo từng token

NVFP4 mạnh mẽ hơn MXFP8, vì vậy chúng tôi áp dụng nó một cách có chọn lọc. Chúng tôi lượng tử hóa các chuyên gia MoE vì chúng chiếm ưu thế về kích thước mô hình và lưu lượng bộ nhớ rollout, trong khi phần còn lại của mô hình vẫn giữ ở định dạng BF16 trừ khi được cấu hình khác.

Ví dụ, DeepSeek-V3 có tổng cộng khoảng 671 tỷ tham số. Các chuyên gia MoE của nó chiếm:

Đó là khoảng 97,8% mô hình. Do đó, việc nhắm mục tiêu vào các chuyên gia MoE sẽ thu được hầu hết lợi ích về bộ nhớ mà không cần ép buộc mọi lớp vào định dạng độ chính xác mạnh nhất.

Tại sao không sử dụng trực tiếp công thức tiền huấn luyện NVFP4?

Công thức tiền huấn luyện NVFP4 gốc được thiết kế cho việc tiền huấn luyện quy mô lớn, nơi mục tiêu là bảo toàn hướng tối ưu hóa thô trên nhiều token trong khi vẫn sử dụng các GEMM FP4. Nó kết hợp các GEMM lớp tuyến tính FP4 với một số bộ ổn định: các lớp được chọn vẫn ở độ chính xác cao hơn, chia tỷ lệ trọng số nhất quán giữa tiến và ngược, và đường dẫn huấn luyện sử dụng làm tròn ngẫu nhiên (stochastic rounding - SR) và Biến đổi Hadamard ngẫu nhiên (Random Hadamard Transforms - RHT). Trong bài báo, SR được áp dụng cho các gradient để giảm độ chệch lượng tử hóa và tạo ra các gradient lượng tử hóa không chệch, trong khi RHT phân tán các giá trị ngoại lai (outliers) cấp khối có độ lớn lớn, đặc biệt là đối với đầu vào GEMM trọng số-gradient.

Đó là một điểm khởi đầu tốt, nhưng RL có một bề mặt lỗi khác:

Công thức RL NVFP4 không kết hợp mọi phần của công thức tiền huấn luyện. Chúng tôi nhắm mục tiêu vào lượng tử hóa trọng số chuyên gia MoE, chia tỷ lệ kích hoạt theo từng token, kiểm soát độ chính xác nhất quán và các GEMM ngược BF16 với các toán hạng gốc hoặc đã giải lượng tử hóa có thể lựa chọn.

Chia tỷ lệ kích hoạt theo từng token

Hệ thống phân cấp NVFP4 hai cấp rất mạnh mẽ, nhưng phạm vi của tỷ lệ kích hoạt FP32 phải được chọn cẩn thận. Như đã thảo luận trong báo cáo kỹ thuật Cursor Composer 2, việc chia tỷ lệ NVFP4 theo từng tensor có thể làm cho quá trình huấn luyện thay đổi theo batch, và việc chia sẻ tỷ lệ giữa các token có thể làm rò rỉ thông tin token tương lai vào các biểu diễn token quá khứ. Nếu một token chia sẻ tỷ lệ của nó với các token khác, biểu diễn lượng tử hóa của nó phụ thuộc vào thành phần batch. Điều này đặc biệt có vấn đề đối với RL, nơi lập lịch rollout và độ dài chuỗi thay đổi.

Do đó, công thức của chúng tôi tính toán một tỷ lệ kích hoạt FP32 cho mỗi token trực tuyến. Điều này khu trú các giá trị ngoại lai kích hoạt vào một token, loại bỏ hiện tượng hiệu chuẩn kích hoạt tĩnh, và cho phép rollout SGLang và huấn luyện Megatron sử dụng cùng một phạm vi tỷ lệ kích hoạt.

Về phía rollout, việc tính toán tỷ lệ FP32 theo từng token được hợp nhất vào đường dẫn nhân lượng tử hóa kích hoạt của FlashInfer: cùng một lệnh gọi phát ra các kích hoạt FP4 đã đóng gói và các tỷ lệ khối E4M3 cũng trả về các tỷ lệ FP32 theo từng token. Kết quả là, việc chia tỷ lệ kích hoạt theo từng token không yêu cầu một bước hiệu chuẩn tỷ lệ riêng biệt.

Sự nhất quán giữa huấn luyện và suy luận cũng đòi hỏi sự khớp về tính song song. Nếu tỷ lệ FP32 được tính toán cho mỗi token trong một phân vùng song song tensor chuyên gia (expert-tensor-parallel), SGLang và Megatron nên sử dụng cùng kích thước ETP. Nếu không, mỗi bên có thể thấy một phân vùng khác nhau của tensor và tính toán một tỷ lệ khác nhau.

Các lớp SwiGLU MoE thêm một hợp đồng quan trọng khác. SGLang và Megatron thường hợp nhất các phép chiếu cổng (gate) và lên (up) thành một GEMM, vì vậy cả hai tensor phải chia sẻ cùng một tỷ lệ FP32 trong quá trình chuyển đổi và cập nhật trọng số trực tiếp ngay cả khi checkpoint Hugging Face lưu trữ chúng riêng biệt. Miles thực thi điều này bằng cách lượng tử hóa các cặp gate/up cùng nhau trong đường dẫn xuất NVFP4.

Chúng tôi đã triển khai và đưa công thức NVFP4 theo từng token lên toàn bộ ngăn xếp: