Tối ưu hóa chi phí chưng cất tri thức (Knowledge Distillation) để triển khai quy mô lớn

Quay lại các bài viết

Chưng cất tri thức (knowledge distillation), quá trình huấn luyện một mô hình học sinh nhỏ hơn để đạt hiệu suất tương đương với mô hình giáo viên lớn hơn, là một kỹ thuật nổi tiếng trong Machine Learning. Với làn sóng các mô hình ngôn ngữ lớn (LLM) mã nguồn mở gần đây như gpt-oss, Qwen, GLM hoặc Kimi, đây đã trở thành một chủ đề nghiên cứu phổ biến trở lại. Việc triển khai các mô hình rất lớn này rất tốn kém: mô hình Kimi-K3 gần đây có 2,8 nghìn tỷ tham số và cần khoảng 3TB VRAM chỉ để tải. Do đó, việc nén chúng thành các mô hình nhỏ hơn và khôi phục các khả năng ban đầu thông qua chưng cất tri thức đã trở thành thông lệ tiêu chuẩn, với các công ty như Nvidia (Nemotron 3 Puzzle 75B) hoặc Multiverse Computing (Hypernova 60B) gần đây đã phát hành các mô hình nén chất lượng cao.

Bước chưng cất là yếu tố quyết định phần lớn chất lượng cuối cùng, nhưng đây cũng thường là phần tốn kém nhất trong quy trình. Việc giữ cả mô hình giáo viên và học sinh trong bộ nhớ, đồng thời tạo ra phân phối xác suất trên toàn bộ từ vựng cho mỗi token, đòi hỏi lượng VRAM khổng lồ, thường chỉ khả thi với hàng trăm GPU và các chiến lược song song hóa tensor (tensor-parallelism) cẩn thận. Bài báo mới nhất của chúng tôi, "Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss", giải quyết vấn đề này bằng hai thay đổi hệ thống: lưu đệm (cache) các logit top-K của giáo viên một lần để giáo viên không bao giờ phải nằm trong bộ nhớ cùng với học sinh, và một hàm mất mát KL-divergence mới tiết kiệm bộ nhớ, giúp tránh việc tạo ra ma trận kích thước từ vựng × độ dài chuỗi đầy đủ, từ đó giảm mức sử dụng VRAM xuống thấp hơn nhiều so với các triển khai mặc định trong các thư viện như PyTorch hoặc NVIDIA Megatron-Bridge. Cùng với nhau, hai thay đổi này cắt giảm chi phí huấn luyện đủ để việc "chữa lành" (healing) ngữ cảnh dài trở nên khả thi trên một GPU duy nhất, và đủ rẻ để thực hiện các thử nghiệm quy mô lớn một cách thực tế.

Tại sao khôi phục bằng chưng cất lại đắt đỏ

Thiết lập tiêu chuẩn, chưng cất trực tuyến (online distillation) sử dụng hàm mất mát Kullback-Leibler (KL loss), giữ cả giáo viên và học sinh trong bộ nhớ cùng một lúc. Tại mỗi bước huấn luyện, giáo viên thực hiện một lượt truyền xuôi (forward pass) đầy đủ để tạo ra phân phối đầu ra của nó, và học sinh được huấn luyện để khớp với phân phối đó. Đây là thiết lập biểu đạt tốt nhất vì phân phối đầy đủ của giáo viên luôn sẵn có, nhưng nó cũng là thiết lập tốn kém nhất về bộ nhớ và tính toán: hai tensor từ vựng đầy đủ phải được lưu giữ cho mỗi vị trí token, và giáo viên phải được tính toán lại trong mỗi bước dù hành vi của nó không thay đổi trong suốt quá trình huấn luyện.

Ví dụ thực tế, gpt-oss-120b có từ vựng gồm 201.088 token. Với độ dài chuỗi 32K và kích thước batch là 4, riêng tensor xác suất của giáo viên đã có hình dạng 4 × 201.088 × 32.768; ở định dạng bfloat16, đó đã là khoảng 50GB VRAM cho một tensor duy nhất. Thêm vào đó là gradient, các giá trị kích hoạt (activations), trọng số mô hình và trạng thái bộ tối ưu hóa (optimizer states), một lần lặp huấn luyện chưng cất có thể đạt đỉnh khoảng 250GB VRAM, vượt quá khả năng cung cấp của cả GPU H200 hoặc B200. Trong bài viết này, chúng tôi chỉ ra rằng việc định dạng lại hàm mất mát KL để xử lý dữ liệu theo từng khối (chunks) giúp giảm chi phí này xuống gần như bằng không.

Hàm KL dày đặc (Dense KL) tăng vọt lên khoảng 250GB, vượt quá dung lượng 141GB của một GPU H200. Hàm mất mát theo khối đã hợp nhất không tạo ra mức tăng vọt đó và đạt đỉnh ở khoảng 128GB. Nguồn: Hình 1 của bài báo.

Hai thay đổi hệ thống

Chưng cất ngoại tuyến (Offline distillation). Thay vì tính toán lại giáo viên ở mỗi bước, chúng tôi tính toán đầu ra của nó một lần, lưu đệm 100 token có khả năng cao nhất cho mỗi vị trí, và huấn luyện học sinh dựa trên bộ đệm đó. Giáo viên không bao giờ phải nằm trong bộ nhớ trong quá trình huấn luyện và không cần phải chạy lại sau khi bộ đệm đã tồn tại, vì vậy cùng một bộ đệm có thể được tái sử dụng cho nhiều thử nghiệm cắt bỏ (ablations).

Hàm mất mát KL hợp nhất, theo khối (fused, chunked KL loss). Để hiểu tại sao bản thân hàm mất mát này lại đắt đỏ, hãy hình dung những gì nó thực sự xây dựng: đối với mỗi vị trí token trong một chuỗi và mỗi từ trong từ vựng, hàm mất mát cần một con số mô tả mức độ sai lệch giữa dự đoán của học sinh và giáo viên. Nếu trình bày dưới dạng lưới, đó là một hàng cho mỗi mục từ vựng và một cột cho mỗi vị trí chuỗi; với từ vựng hơn 100K từ và chuỗi dài, lưới đó là khổng lồ, và cách tính hàm mất mát KL mặc định sẽ xây dựng toàn bộ lưới đó trước khi có thể tạo ra một con số duy nhất.

Chúng tôi so sánh ba cách tính cùng một hàm mất mát này, tất cả đều tương đương về mặt toán học:

GIF dưới đây cho thấy sự khác biệt giữa phương pháp dày đặc (dense) và phương pháp hợp nhất theo khối (fused-chunked): một phương pháp xây dựng toàn bộ lưới so sánh và giữ lại tất cả, phương pháp kia xây dựng và loại bỏ từng lát cắt một, vì vậy bộ nhớ không bao giờ vượt quá một khối duy nhất.

Chúng tôi đã mở mã nguồn triển khai hàm mất mát theo khối: github.com/CompactifAI/Full-Chunked-KL-Loss

Điều này thay đổi những gì trong thực tế

Bảng dưới đây so sánh trực tiếp cả bốn thiết lập: chưng cất trực tuyến và ba triển khai hàm mất mát ngoại tuyến vừa mô tả. Khi so sánh trên một GPU H200 duy nhất với Llama 3.1 8B Instruct làm giáo viên và mô hình Llama 3.2B làm học sinh ở ngữ cảnh 8K token, cả bốn đều đạt được hàm mất mát huấn luyện gần như giống hệt nhau, mặc dù các lần chạy ngoại tuyến chỉ huấn luyện dựa trên các logit top-100 được lưu đệm cho mỗi token.

Các đường cong mất mát chồng lên nhau gần như chính xác trên cả bốn phương pháp, xác nhận rằng chưng cất ngoại tuyến với logit top-100 được lưu đệm không bị mất mát thông tin so với chưng cất trực tuyến. Nguồn: Hình 2 của bài báo. Ở độ dài chuỗi này, hàm mất mát hợp nhất theo khối chưa phải là lựa chọn nhanh nhất, việc tính toán thêm phép chiếu ngược (backward-pass projection) làm giảm một chút tốc độ, nhưng lợi thế thực sự của nó chỉ xuất hiện khi độ dài ngữ cảnh tăng lên, điều mà phần tiếp theo sẽ chứng minh.

Mở rộng quy mô đến độ dài ngữ cảnh dài

Để thấy rõ hơn mô hình mở rộng quy mô, chúng tôi đã chạy một điểm chuẩn cô lập trên một mạng lưới chiếu đầu ra đơn giản (không có thân transformer, chỉ có nhân hàm mất mát). Ở 32K token, bộ nhớ đỉnh giảm từ 85,2 GiB với hàm mất mát dày đặc xuống còn 5,45 GiB với phiên bản phân khối hoàn toàn, giảm 15,6 lần, và hàm mất mát dày đặc thất bại hoàn toàn từ 64K token trở đi. Ở 256K token, hàm mất mát phân khối hoàn toàn sử dụng 11,6 GiB so với 134,2 GiB của biến thể phân khối tốt thứ hai, và nhanh hơn khoảng 3,3 lần mỗi lần lặp ở độ dài đó.

Khi chưng cất mô hình GPT-OSS 20B ở ngữ cảnh 32.768 token, bộ nhớ được giải phóng bởi hàm mất mát hợp nhất cho phép thiết lập thu nhỏ từ bốn node GPU xuống còn một. Thời gian mỗi bước giảm từ 57,0 xuống 12,23 giây, nhanh hơn khoảng 5 lần, và thông lượng trên mỗi GPU tăng từ 74,2 lên 345,7 TFLOP/s.

Mô hình học sinh thu được

Thiết lập ngoại tuyến hiệu quả là yếu tố giúp chiến dịch chưng cất quy mô lớn trở nên hợp túi tiền ngay từ đầu. Mô hình học sinh nhỏ gọn thu được, được chưng cất từ Llama 3.1 8B Instruct xuống còn khoảng 3,2 tỷ tham số, giữ lại hầu hết độ chính xác của giáo viên trên BoolQ và HellaSwag, duy trì khoảng cách trong vòng chín điểm trên MMLU, với số lượng tham số chưa bằng một nửa.

Học sinh giữ lại hầu hết độ chính xác ngữ cảnh ngắn của giáo viên với kích thước chưa bằng một nửa. Nguồn: Hình 6 của bài báo.

Công trình này là một phần trong nghiên cứu đang diễn ra của Multiverse Computing nhằm làm cho việc chưng cất và "chữa lành" trở nên thiết thực để chạy ở quy mô lớn, không chỉ như một công thức dùng một lần, mà là thứ mà các nhóm có thể lặp lại với chi phí thấp. Bài báo cũng đề cập đến các thử nghiệm cắt bỏ bổ sung, chẳng hạn như cách lựa chọn hàm mất mát và đóng gói chuỗi ảnh hưởng đến chất lượng khôi phục.

Bạn muốn biết chi tiết kỹ thuật đầy đủ, bao gồm gradient dạng đóng (closed-form gradient) đằng sau hàm mất mát hợp nhất theo khối và cấu hình huấn luyện hoàn chỉnh? Hãy đọc toàn bộ bài báo, hoặc liên hệ với nhóm của chúng tôi để thảo luận về việc áp dụng điều này vào các quy trình chưng cất của riêng bạn.

Chúng tôi cũng đã mở mã nguồn triển khai hàm mất mát theo khối: github.com/CompactifAI/Full-Chunked-KL-Loss