Speculative Decoding: Kỹ thuật tăng tốc LLM lên 3 lần
Speculative decoding là kỹ thuật tăng tốc sinh văn bản của mô hình ngôn ngữ lớn (LLM) bằng cách dùng một mô hình nhỏ dự đoán trước nhiều token, sau đó mô hình lớn xác nhận song song. Phương pháp này giúp giảm đáng kể độ trễ so với giải mã tự hồi quy (autoregressive) truyền thống mà không làm thay đổi kết quả đầu ra. Bài viết này phân tích cơ chế hoạt động, lợi ích và cách triển khai speculative decoding trong thực tế.
Nguồn tham khảo: Hugging Face Blog: Assisted Generation, arXiv:2211.17192
Vì sao sinh văn bản LLM chậm?
Quá trình sinh văn bản của LLM hoạt động theo từng token một: mô hình nhận toàn bộ chuỗi đã sinh và chạy một lượt forward pass để dự đoán token tiếp theo, sau đó token mới được thêm vào và lặp lại. Với mô hình lớn, mỗi forward pass tốn nhiều thời gian vì phải tải toàn bộ trọng số qua bộ nhớ, dẫn đến độ trễ cao khi sinh hàng trăm token liên tiếp.
Điểm nghẽn thực sự không phải là phép tính mà là băng thông bộ nhớ (memory bandwidth): việc nạp trọng số mô hình từ GPU memory vào lõi tính toán chiếm phần lớn thời gian của forward pass. Chính vì vậy, nếu có thể dự đoán được nhiều token trong một lần, ta sẽ giảm được số lần forward pass cần thiết và tận dụng tốt hơn băng thông bộ nhớ.
Speculative decoding hoạt động thế nào?
Ý tưởng cốt lõi của speculative decoding là dùng một mô hình nhỏ (draft model) để đề xuất một chuỗi token dự đoán, sau đó mô hình lớn (target model) xác nhận toàn bộ chuỗi trong một forward pass duy nhất. Quy trình gồm các bước:
- Draft model dự đoán: Mô hình nhỏ, nhanh sinh ra K token tiếp theo dựa trên chuỗi hiện tại. Vì mô hình nhỏ có chi phí thấp, việc sinh K token này rất nhanh.
- Target model xác nhận: Mô hình lớn chạy một forward pass trên chuỗi gồm K token được đề xuất, tính xác suất cho từng vị trí.
- Chấp nhận hoặc sửa: Nếu token dự đoán của mô hình nhỏ khớp với token có xác suất cao nhất của mô hình lớn, token đó được chấp nhận ngay. Nếu không khớp, token từ mô hình lớn được dùng và chuỗi dừng lại tại điểm đó.
- Lặp lại: Quá trình này lặp lại cho đến khi sinh đủ số token mong muốn.
Với cách làm này, mô hình lớn chỉ cần một forward pass để xác nhận K token thay vì K forward pass riêng lẻ. Trong thực tế, với draft model tốt, tỷ lệ chấp nhận có thể đạt đến 70-80%, giúp tăng tốc tổng thể từ 2 đến 3 lần trên phần cứng thông thường.
Các biến thể và cải tiến
Bên cạnh phiên bản gốc, nhiều biến thể của speculative decoding đã được phát triển để tăng hiệu quả:
- Assisted generation: Triển khai speculative decoding trong thư viện Hugging Face Transformers, sử dụng draft model bất kỳ và tự động đồng bộ phân phối token.
- Speculative sampling: Cải tiến thuật toán chấp nhận token dựa trên xác suất mẫu, đảm bảo phân phối đầu ra không đổi so với target model gốc.
- Medusa / EAGLE: Dùng các head dự đoán bổ sung (draft heads) đặt trên chính mô hình lớn thay vì dùng một mô hình nhỏ riêng biệt, giảm chi phí triển khai và tăng tốc nhẹ.
- Self-speculative decoding: Chính mô hình lớn tự sinh token dự đoán bằng cách bỏ qua một số lớp (layer) ở chế độ dự thảo, không cần mô hình phụ trợ.
Những biến thể này giúp speculative decoding áp dụng được cho nhiều loại mô hình và bài toán khác nhau, từ code completion tới chatbot thời gian thực.
Lợi ích và hạn chế
Lợi ích chính:
- Giảm độ trễ đáng kể cho ứng dụng real-time như chatbot, code completion, dịch máy
- Kết quả đầu ra giữ nguyên so với greedy decoding hoặc sampling thông thường
- Không cần thay đổi kiến trúc mô hình, chỉ cần thêm draft model hoặc draft head
- Tận dụng tốt phần cứng có sẵn mà không cần GPU mới
Hạn chế:
- Cần thêm bộ nhớ cho draft model hoặc draft head
- Hiệu quả phụ thuộc vào độ khớp giữa draft model và target model
- Trong trường hợp token khó dự đoán, tỷ lệ chấp nhận thấp và lợi ích giảm
- Độ phức tạp triển khai tăng khi kết hợp với batching và serving
Triển khai speculative decoding với Hugging Face
Thư viện Transformers của Hugging Face hỗ trợ assisted generation trực tiếp qua tham số assistant_model khi gọi generate():
from transformers import AutoModelForCausalLM, AutoTokenizer
target = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3-8B")
draft = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3-8B")
outputs = target.generate(
**tokenizer("Viết về Việt Nam", return_tensors="pt"),
assistant_model=draft,
max_new_tokens=200
)Chỉ cần đặt tham số assistant_model, Transformers sẽ tự động chạy cơ chế chấp nhận token, giúp tăng tốc đáng kể mà không cần sửa code phức tạp.
Khi nào nên dùng speculative decoding?
Speculative decoding phù hợp nhất khi ứng dụng yêu cầu độ trễ thấp và có draft model đủ tốt, chẳng hạn code completion trong IDE, chatbot customer support, hoặc suy luận nhúng trên thiết bị biên. Ngược lại, nếu chỉ quan tâm throughput tổng và batch rất lớn, các kỹ thuật khác như continuous batching hoặc quantization có thể hiệu quả hơn.
Trong tương lai, speculative decoding sẽ ngày càng phổ biến trong các inference serving framework như vLLM, TGI và TensorRT-LLM vì lợi ích rõ ràng về độ trễ mà không đánh đổi chất lượng đầu ra. Đây là một trong những kỹ thuật quan trọng giúp LLM trở nên hữu dụng hơn trong các tình huống thời gian thực.
