NVIDIA ra mắt Nemotron-Labs-3-Puzzle-75B: Mô hình MoE nén giúp tăng gấp đôi hiệu suất máy chủ

Các mô hình MoE lai quy mô lớn như Nemotron-3-Super rất chính xác nhưng chi phí vận hành lại đắt đỏ. Các tham số hoạt động (active parameters), bộ nhớ đệm KV (KV cache) và trạng thái Mamba (Mamba state) giới hạn số lượng người dùng mà một node có thể xử lý ở một tốc độ token nhất định trên mỗi người dùng. Nhóm NVIDIA AI đã phát hành Nemotron-Labs-3-Puzzle-75B-A9B, một biến thể nén của Nemotron-3-Super. Mô hình gốc có tổng cộng 120,7 tỷ tham số, trong đó có 12,8 tỷ tham số hoạt động. Mô hình nén có tổng cộng 75,3 tỷ tham số, với 9,3 tỷ tham số hoạt động.

Mục tiêu triển khai đã được ấn định trước khi quá trình tìm kiếm kiến trúc bắt đầu. Mục tiêu thứ nhất là đạt thông lượng máy chủ gấp 2 lần ở mức 100 token mỗi giây cho mỗi người dùng. Mục tiêu thứ hai là xử lý 8 yêu cầu 1 triệu token đồng thời trên một card H100 duy nhất. Ba checkpoint trên Hugging Face bao gồm: BF16, FP8 và NVFP4.

Tóm tắt

Nemotron-3-Super là một mô hình MoE lai giữa Mamba và Transformer. Puzzle-75B-A9B bảo toàn chính xác bố cục khối của mô hình gốc. Nó có 88 khối: 40 khối Mamba, 40 khối MoE và 8 khối attention.

Những thay đổi nằm ở dung lượng bên trong các khối đó:

Số lượng chuyên gia được định tuyến (routed experts), kích thước chuyên gia chia sẻ (shared expert size) và kích thước tiềm ẩn (latent size) của MoE không thay đổi. Các lớp attention được giữ nguyên. Lý do được nêu trong nghiên cứu là vì Nemotron-3-Super vốn đã rất hiệu quả về bộ nhớ đệm KV. Các lớp Mamba được cắt tỉa đồng nhất vì các framework suy luận không hỗ trợ kích thước trạng thái SSM khác nhau cho mỗi lớp.

Kết quả không phải là một mô hình giáo viên (teacher) được thu nhỏ đồng nhất. Hình trên cho thấy sự phân bổ theo chiều sâu. Puzzle bảo toàn dung lượng ở các lớp giữa và lớp cuối được chọn, đồng thời cắt giảm mạnh ở những nơi khác.

Điểm chuẩn và Hiệu suất

Bảng dưới đây báo cáo thông lượng tổng thể tối ưu Pareto trên một node 8xB200 duy nhất, với giải mã đơn bước (single-step decoding).

Cả hai mô hình đều được phục vụ với trọng số NVFP4, bộ nhớ đệm KV FP8 và trạng thái Mamba FP16 tương ứng. Do đó, khoảng cách hiệu suất phản ánh quá trình nén chứ không phải do thay đổi định dạng số. Chế độ 50K/2K (nặng về prefill) đạt mức tăng thấp nhất. Chế độ 8K/64K (nặng về decode) đạt mức tăng cao nhất.

Trên một node 8xH100 duy nhất tại UT = 100, mức tăng nhỏ hơn. Cụ thể là 1,91 lần trên 50K/2K và 1,82 lần trên 8K/64K. Cả hai mô hình ở đây đều sử dụng trọng số FP8, bộ nhớ đệm KV FP8 và trạng thái Mamba FP32.

Trên một card H100 duy nhất với ngữ cảnh 1 triệu token, ràng buộc bị chuyển từ tính toán sang bộ nhớ. Trọng số NVFP4 của Super chiếm khoảng 70 GB trong tổng ngân sách 80 GB HBM. Mỗi yêu cầu 1 triệu token tiêu tốn khoảng 4 GB bộ nhớ đệm KV. Do đó, khả năng xử lý đồng thời hiệu dụng là 1.

Trọng số NVFP4 của Puzzle-75B-A9B chiếm khoảng 44,5 GB. Bố cục attention không thay đổi nên chi phí KV cho mỗi yêu cầu là như cũ. Khả năng xử lý đồng thời ở mức 1 triệu token tăng lên 8. Thông lượng giải mã tổng hợp ở mức đồng thời đó gấp khoảng 4 lần thông lượng yêu cầu đơn lẻ của Super. Tốc độ prefill cho một prompt 990K token nhanh hơn khoảng 1,2 lần.

Cách thức hoạt động của Iterative Puzzle

Puzzle là một framework tìm kiếm kiến trúc thần kinh phân rã, được triển khai ở đây dưới tên gọi Puzzletron. Nó xác định một không gian tìm kiếm rời rạc cho các cách triển khai lớp thay thế. Mỗi phương án thay thế đều nhận được một điểm chất lượng. Sau đó, một chương trình số nguyên hỗn hợp sẽ chọn một phương án cho mỗi lớp dựa trên ràng buộc triển khai.

Ba kỹ thuật cắt tỉa tạo nên không gian tìm kiếm:

Kết quả SSM được đo lường. Việc giảm 128 kênh xuống 96 giúp tăng tốc nhân SSM từ 1,2 đến 1,3 lần trong quá trình giải mã. Điều này duy trì ở các kích thước batch từ 8 đến 512. Các kênh được xếp hạng theo đóng góp ước tính vào đầu ra của lớp Mamba. Ước tính này được lấy trung bình trên 67 triệu token dữ liệu kiểm chứng. Phụ lục A cho thấy phương pháp này vượt trội hơn so với việc chọn kênh ngẫu nhiên khi cắt tỉa mạnh.

Công thức ban đầu giả định rằng tác động chất lượng của việc thay thế là cộng dồn xấp xỉ. Mỗi khối ứng viên được chấm điểm bên trong mô hình gốc chưa sửa đổi. Điều này bỏ qua các tương tác bậc cao giữa các thay thế.

Iterative Puzzle luân phiên giữa nén có giới hạn và phục hồi bằng chưng cất tri thức (knowledge distillation) ngắn hạn. Nó xây dựng một chuỗi M0, M1, … MR thay vì nhảy thẳng đến mục tiêu. Điểm số được tính toán lại dựa trên mô hình nén hiện tại, thay vì mô hình gốc ban đầu.

Ba giai đoạn đã được sử dụng:

Bảng trên so sánh điều này với đường cơ sở Puzzle đơn bước tại cùng một mục tiêu. Quy trình ba bước đạt trung bình 69,05 trên mười điểm chuẩn, so với 68,48. Mức tăng xuất hiện trên MMLU-Pro, GPQA, HLE, AA-LCR, LiveCodeBench, SciCode và RULER-256K. IFBench-Instruction giảm 0,2 điểm và IFBench-Prompt giảm 0,5 điểm.

Phục hồi: Chưng cất, RL và Độ dài văn bản (Verbosity)

Chưng cất tri thức được thực hiện trên 30% dữ liệu tiền huấn luyện và 70% dữ liệu SFT từ Nemotron-3-Nano. Trong giai đoạn Puzzle, KD sử dụng độ dài chuỗi 32K. Sau đó, quá trình phục hồi được huấn luyện ở mức 128K và mở rộng lên 512K. Ngân sách lên tới 100 tỷ token, với batch toàn cục 16 triệu token, trong Megatron-LM.

Hậu huấn luyện RL áp dụng Giai đoạn 2 của quy trình RL Nemotron-3-Super, tập trung vào kỹ thuật phần mềm. Giai đoạn 2.1 thực hiện so sánh sử dụng công cụ đơn bước. Giai đoạn 2.2 chuyển sang RL sandbox đầu-cuối, nơi các tác nhân chạy tới 200 lượt. Cả hai giai đoạn đều sử dụng mức phạt KL bằng 0. Nhóm đã quét qua các tốc độ học (learning rates), sau đó lấy trung bình các trọng số thu được.

Hình 4 ở trên cho thấy đóng góp của từng giai đoạn. KD ngữ cảnh ngắn phục hồi hầu hết các danh mục lên hơn 97% so với Nemotron-3-Super. KD ngữ cảnh dài sau đó nâng cao cụ thể các điểm chuẩn đầu vào dài và thế hệ dài. Nhóm nghiên cứu cho biết tác động của RL trong các thí nghiệm này là nhỏ.

Độ dài văn bản là một chi tiết thầm lặng. Sau lần lặp Puzzle cuối cùng, mô hình tạo ra 132% số lượng token của Super. Con số đó giảm xuống còn 99% sau quy trình phục hồi đầy đủ.

Triển khai: Lượng tử hóa và Dự đoán đa token (Multi-Token Prediction)

Hai công thức lượng tử hóa hậu huấn luyện đã được tạo ra: FP8 W8A8 nhắm vào Hopper và NVFP4 W4A4 nhắm vào Blackwell.

Cả hai công thức đều được hiệu chuẩn trên 256 mẫu SFT hậu huấn luyện. NVFP4 sử dụng hiệu chuẩn tối đa, không phải tìm kiếm độ nhạy AutoQuantize như đã dùng cho Super. Checkpoint thu được bị lượng tử hóa mạnh hơn một chút và có hiệu suất tương đương.

NVFP4 không được hỗ trợ nguyên bản trên Hopper. Nó vẫn được sử dụng cho mục tiêu H100 ngữ cảnh 1 triệu token vì dung lượng HBM bị giới hạn ở đó.

Puzzle-75B-A9B kế thừa đầu MTP chia sẻ từ Super. Các tham số được chia sẻ qua các bước MTP, vì vậy một đầu áp dụng đệ quy khi suy luận. Việc chuyển trực tiếp đầu đã huấn luyện của Super cho độ dài chấp nhận tương tự.

Nhóm nghiên cứu sau đó xác định sự không khớp giữa huấn luyện và suy luận. Huấn luyện MTP có sự ép buộc của giáo viên (teacher-forced) cung cấp toàn bộ chuỗi trạng thái ẩn đã dịch chuyển. Ngược lại, dự thảo tự hồi quy (autoregressive drafting) cung cấp hỗn hợp các trạng thái ẩn từ mô hình mục tiêu và trạng thái ẩn do MTP tạo ra. Tỷ lệ chấp nhận giảm ở các vị trí dự thảo sâu hơn.

Việc tiếp tục huấn luyện trên đầu đã chuyển đổi giải quyết vấn đề này. Trên SPEED-Bench ở độ dài dự thảo 7, độ dài chấp nhận trung bình tăng từ 3,45 lên 4,34. Đó là khoảng 25% đến 30%, tập trung ở các vị trí dự thảo sau này. Không giống như Super, checkpoint NVFP4 hầu như không bị suy giảm: 4,31 so với 4,34.

Nén giúp ích ở đâu và gây hại ở đâu

Tóm tắt của chính bài báo nghiên cứu là các đánh giá về khả năng tuân thủ hướng dẫn và tác nhân bị ảnh hưởng nhiều nhất. Arena-Hard-V2 là trường hợp tệ nhất, với -4,2 điểm. RULER duy trì trong khoảng 1 đến 2 điểm ở mức 256K, 512K và 1M.

Ba kết quả BF16 không bị thoái lui. AA-LCR tăng 0,1, Scale AI Multi-Challenge hòa ở mức 56,6 và TauBench Telecom tăng 0,4.