Tencent ra mắt AngelSpec: Khung huấn luyện thống nhất cho mô hình suy luận dự đoán

Tencent đã phát hành AngelSpec, một framework huấn luyện mã nguồn mở, native-torch dành cho các mô hình dự thảo (draft models) trong giải mã suy đoán (speculative decoding). Bản phát hành này bao gồm cả dự đoán đa token tự hồi quy (autoregressive multi-token prediction - MTP) và dòng mô hình block-parallel DFlash.

Hầu hết các nghiên cứu về giải mã suy đoán đều tìm kiếm một mô hình dự thảo (drafter) hoạt động tốt trên một hỗn hợp benchmark trung bình. Tuy nhiên, lưu lượng truy cập thực tế khi phục vụ (serving traffic) không giống với hỗn hợp đó. AngelSpec coi tính không đồng nhất của khối lượng công việc là một ràng buộc thiết kế hàng đầu, đồng thời chuyên biệt hóa cấu trúc, dữ liệu huấn luyện và độ sâu xác thực xung quanh yếu tố này.

Tại sao một mô hình dự thảo phổ quát lại hoạt động kém hiệu quả

Giải mã suy đoán là phương pháp không mất mát dữ liệu (lossless). Một mô hình dự thảo nhẹ sẽ đề xuất một vài token tương lai, và mô hình mục tiêu sẽ xác thực chúng cùng lúc trong một lượt truyền xuôi (forward pass) bằng cách sử dụng lấy mẫu loại trừ (rejection sampling). Do đó, khả năng tăng tốc phụ thuộc vào hai yếu tố: bao nhiêu token dự thảo được chấp nhận và thời gian hoàn thành vòng lặp dự thảo–xác thực mất bao lâu.

Hai đại lượng này biến thiên theo hướng ngược nhau tùy theo từng lĩnh vực. Trong các cuộc hội thoại mở có độ entropy cao, nhiều phần tiếp nối đều hợp lệ về mặt ngữ nghĩa. Mô hình mục tiêu có thể chọn bất kỳ phần nào trong số đó, vì vậy tỷ lệ chấp nhận giảm nhanh theo độ sâu của đề xuất. Việc tạo và xác thực một khối dài gây lãng phí tài nguyên tính toán. Phương pháp dự thảo MTP tự hồi quy phù hợp với chế độ này vì nó đề xuất một chuỗi ứng viên ngắn hơn.

Mã nguồn và suy luận toán học hoạt động theo cách khác. Cú pháp lập trình, các định danh lặp lại, biểu thức hình thức và các bước suy luận từng bước ràng buộc các token tương lai chặt chẽ hơn. Những khối lượng công việc này tạo ra các đoạn có thể dự đoán dài hơn, đây chính xác là những gì mà phương pháp dự thảo song song theo khối (block-parallel drafting) tối ưu hóa rất tốt.

Do đó, AngelSpec cung cấp hai mô hình dự thảo bổ trợ cho nhau thay vì một giải pháp thỏa hiệp duy nhất. Mô hình MTP được huấn luyện trên dữ liệu phong phú, đa dạng hướng tới hội thoại. Mô hình block-diffusion được tăng cường với các mẫu tập trung vào mã nguồn và toán học.

Đường dẫn MTP: Kiểm thử trong thời gian huấn luyện (Training-Time Test) và triển khai mô hình mục tiêu

Mô hình Hy3 gốc được huấn luyện với một lớp MTP duy nhất và không có cơ chế unrolling tự điều kiện đệ quy. Khi suy luận, khối này có thể được tái sử dụng theo kiểu đệ quy, nhưng mục tiêu huấn luyện chưa bao giờ chuẩn bị cho nó một chuỗi tự tạo dài. Các lỗi tích tụ theo độ sâu, vì vậy vị trí dự thảo thứ hai và thứ ba có tỷ lệ chấp nhận thấp hơn đáng kể so với vị trí đầu tiên.

AngelSpec giải quyết sự không khớp giữa huấn luyện và suy luận này bằng một lược đồ đa độ sâu, chia sẻ tham số. Nó giữ lại D độ sâu dự đoán logic nhưng tái sử dụng một khối MTP vật lý. Trong quá trình huấn luyện, khối đó được unrolling tự hồi quy trong D bước, với mỗi dự đoán được đưa vào lần gọi tiếp theo. Theo nguyên tắc Kiểm thử trong thời gian huấn luyện của EAGLE-3, độ sâu k+1 nhận dự đoán arg max từ độ sâu k thay vì token thực tế (ground-truth). Các tham số được chia sẻ, nhưng sự giám sát vẫn mang tính đặc thù theo độ sâu: mục tiêu giáo viên (teacher target) tiến thêm một vị trí tương lai tại mỗi độ sâu.

Hai lựa chọn khác mang lại phần lớn hiệu quả. Thứ nhất, backbone mục tiêu và head mô hình ngôn ngữ mục tiêu được đóng băng (frozen), và các đầu vào MTP từ backbone được tách rời. Mô hình dự thảo cải thiện phân phối đề xuất của nó mà không làm thay đổi phân phối mà quá trình xác thực phải bảo toàn. Thứ hai, quá trình huấn luyện sử dụng triển khai mô hình mục tiêu (target-model rollout): các phản hồi được tạo bởi mô hình mục tiêu đã đóng băng thay vì lấy từ kho ngữ liệu tham chiếu gốc. Điều đó tạo ra các lựa chọn token, quỹ đạo trạng thái ẩn và các mô hình không chắc chắn cục bộ chính xác mà MTP phải xấp xỉ tại thời điểm phục vụ.

Hiệu quả đo lường được tập trung đúng nơi cần thiết. Tại T = 0, tỷ lệ chấp nhận trung bình tăng từ 52,8% lên 66,4%, và độ dài chấp nhận trung bình từ 2,58 lên 2,99. Tỷ lệ ở vị trí đầu tiên gần như được giữ nguyên ở mức 0,799 → 0,814. Các vị trí sâu hơn mang lại sự khác biệt: p3 tăng từ 0,290 lên 0,706 trên GSM8K, và từ 0,387 lên 0,757 trên HumanEval.

DFly: điều kiện mục tiêu lai và head tự hồi quy có điều kiện tiền nhiệm

DFly là kiến trúc block-diffusion. Nó được xây dựng dựa trên DFlash với hai thay đổi về cấu trúc.

Backbone điều kiện mục tiêu lai: DFlash nối các trạng thái ẩn từ nhiều lớp mục tiêu và biến đổi chúng bằng một lớp kết nối đầy đủ (fully connected layer), tạo ra một đặc trưng ngữ cảnh chung. Ngữ cảnh đơn lẻ đó sau đó được cung cấp cho mọi lớp dự thảo, điều này hạn chế sự chuyên biệt hóa của lớp. Thay vào đó, DFlare học trọng số hợp nhất theo từng lớp, mang lại cho mỗi lớp dự thảo cái nhìn riêng về hệ thống phân cấp mục tiêu, nhưng loại bỏ phép biến đổi xuyên lớp đã học của DFlash. DFly kết hợp chúng: nhánh FC thiết lập một cơ sở ngữ nghĩa chung, và việc hợp nhất theo từng lớp được áp dụng như một sự tinh chỉnh thặng dư (residual refinement) ở phía trên. Nhánh bổ sung chỉ giới thiệu D × T trọng số vô hướng, và các hệ số softmax của nó có thể được tính toán trước sau khi huấn luyện.

Head tự hồi quy có điều kiện tiền nhiệm: Một backbone song song dự đoán từng vị trí khối chỉ từ ngữ cảnh đã được chấp nhận. Nó không thể thấy phần tiếp nối nào thực sự được chọn ở các vị trí dự thảo trước đó, điều này tạo ra sự suy giảm chấp nhận hậu tố. DFly đặt một head tuần tự nhỏ sau backbone song song, chuyển đổi các dự đoán biên theo vị trí thành các phân phối có điều kiện tiền tố. Backbone đắt đỏ vẫn giữ nguyên tính song song hoàn toàn; chỉ có head nhỏ chạy từ trái sang phải.

Hiệu suất ban đầu

Trên Qwen3-8B, DFly đạt độ dài chấp nhận trung bình là 5,41, so với 5,32 của DSpark, 4,57 của DFlash và 3,24 của MTP. Nó đạt kết quả tốt nhất trên cả năm benchmark toán học và mã nguồn. DSpark vẫn dẫn trước một chút trên MT-Bench với 3,77 so với 3,67, điều này nhất quán với việc DFly được định vị cho mã nguồn và toán học.

Trên Hy3-A21B, biên độ cách biệt rộng hơn. DFly đạt 4,79 so với 3,69 của DFlash và 3,00 của MTP — mức tăng tương đối lần lượt là 29,8% và 59,7%. Nó cải thiện tất cả sáu benchmark được báo cáo.

Phân tích tích lũy (ablation) trên Hy3-A21B dưới giải mã tham lam (greedy decoding) cho thấy nguồn gốc của hiệu quả: backbone DFlash 3,77, backbone DFly 4,40, cộng với head Markov 4,56, hiệu chỉnh ẩn 4,60 và dữ liệu mã/toán 4,75. Việc mở rộng dữ liệu bổ sung 700K prompt — 500K mã nguồn từ OpenCodeInstruct và OpenCodeReasoning, 200K toán học từ Big-Math. Bất kỳ prompt nào chia sẻ một đoạn 16-token liên tục với một ví dụ đánh giá đều bị loại bỏ trước khi tạo phản hồi.

Bên trong framework

AngelSpec được xây dựng trên TorchSpec và mở rộng nó ở một số điểm. Nền tảng được phân tách: các engine suy luận chạy mô hình mục tiêu đã đóng băng và truyền các trạng thái ẩn thông qua kho lưu trữ RDMA hỗ trợ bởi Mooncake trực tiếp đến các worker huấn luyện phân tán. Các trạng thái ẩn được thu thập bên trong các tiến trình worker vLLM thông qua các API công khai của vLLM — một hook trích xuất trạng thái ẩn suy đoán và một bộ kết nối KV tùy chỉnh — mà không cần fork engine.

Các phần mở rộng quan trọng cho công trình này:

Các backend được phân tầng: vLLM là ưu tiên hàng đầu, với SGLang và HuggingFace Transformers được hỗ trợ ở tầng cộng đồng.

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

Hãy xem Bài báo, Kho lưu trữ GitHub, Tài liệu, Bộ sưu tập Hugging Face và Bộ sưu tập ModelScope. Mọi ghi nhận cho nghiên cứu này thuộc về các nhà nghiên cứu của dự án.

Michal Sutter là một chuyên gia khoa học dữ liệu với bằng Thạc sĩ Khoa học Dữ liệu từ Đại học Padova. Với nền tảng vững chắc về phân tích thống kê, học máy và kỹ thuật dữ liệu, Michal xuất sắc trong việc chuyển đổi các tập dữ liệu phức tạp thành những thông tin chi tiết có thể hành động.