Giải mã có hỗ trợ

Giải mã có hỗ trợ (Assisted decoding) tăng tốc độ sinh văn bản bằng cách cho phép một mô hình phụ đề xuất các token ứng viên trước khi mô hình chính xác nhận chúng. Mô hình chính xác minh các token ứng viên trong một lần truyền tiến duy nhất. Mô hình phụ nhanh và rẻ, có thể thay thế hàng chục lần truyền tiến đắt đỏ hơn của mô hình chính.

Hướng dẫn này bao gồm các phương pháp giải mã có hỗ trợ trong Transformers.

Giải mã suy đoán

Speculative decoding sử dụng một mô hình trợ lý nhỏ hơn để phác thảo các token ứng viên. Mô hình chính kiểm tra các token này trong một lần chạy. Các token được xác thực sẽ đi vào đầu ra cuối cùng, còn các token bị từ chối sẽ kích hoạt quá trình lấy mẫu tiêu chuẩn. Việc tạo ra nhanh hơn vì mô hình chính thực hiện ít lần chạy forward tốn kém hơn.

Phương pháp này hoạt động tốt nhất khi mô hình trợ lý nhỏ hơn đáng kể so với mô hình chính và sử dụng cùng một tokenizer. Speculative decoding hỗ trợ tìm kiếm tham lam (greedy search) và lấy mẫu (sampling) nhưng không hỗ trợ các đầu vào dạng batch.

Truyền assistant_model vào generate(). Đặt do_sample=True để lấy mẫu lại nếu việc xác thực token thất bại.

tìm kiếm tham lam

sampling

Giải mã tra cứu prompt

Prompt lookup decoding không cần một mô hình trợ lý. Phương pháp này tìm các n-gram chồng lấp trong prompt để đề xuất các token ứng viên. Nếu không có sự khớp nào, nó sẽ quay lại giải mã tự hồi quy thông thường. Cách tiếp cận này phù hợp với các tác vụ dựa trên đầu vào như tóm tắt và dịch, vì các token ứng viên thường phản ánh các mẫu cục bộ trong văn bản nguồn.

Truyền prompt_lookup_num_tokens vào generate(). Điều này xác định số lượng token mà thuật toán sẽ cố gắng sao chép từ phần trước của prompt khi phát hiện một mẫu lặp lại.

giải mã tham lam

sampling

Giải mã suy đoán tự động

Self-speculative decoding sử dụng các lớp trung gian của mô hình làm trợ lý để đề xuất các token ứng viên. Nếu đề xuất khớp, mô hình sẽ thoát sớm và các lớp còn lại sẽ xác minh hoặc sửa các token đó.

Vì tất cả đều là một mô hình duy nhất, trọng số và bộ nhớ đệm được chia sẻ, giúp tăng tốc độ mà không làm tăng thêm chi phí bộ nhớ. Kỹ thuật này chỉ hoạt động với các mô hình được huấn luyện để hỗ trợ logits thoát sớm từ các lớp trung gian.

Truyền assistant_early_exit vào generate() để đặt lớp thoát.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("facebook/layerskip-llama3.2-1B")
model = AutoModelForCausalLM.from_pretrained("facebook/layerskip-llama3.2-1B", dtype="auto")
inputs = tokenizer("Hugging Face is an open-source company", return_tensors="pt")

outputs = model.generate(**inputs, assistant_early_exit=4, do_sample=False, max_new_tokens=20)
tokenizer.batch_decode(outputs, skip_special_tokens=True)

Tầng dự đoán đa token (MTP)

Tầng dự đoán đa token (MTP) phác thảo các token ứng viên bằng các tầng dự đoán bổ sung được lấy từ checkpoint của chính mô hình chính. Các tầng MTP tái sử dụng các embedding và đầu ra (output head) của mô hình chính. Mỗi tầng MTP phác thảo một token, do đó số lượng ứng viên ở mỗi bước khớp với số lượng tầng MTP trong checkpoint. Mô hình chính xác minh các bản phác thảo trong một lần truyền tiến (forward pass).

MTP chỉ hoạt động với các checkpoint được huấn luyện bằng các lớp MTP, chẳng hạn như DeepSeek-V3 và GLM-4.5. Các checkpoint này chứa trọng số MTP và đặt num_mtp_layers trong cấu hình của chúng, được liệt kê dưới dạng num_nextn_predict_layers trong config.json.

Truyền use_mtp=True vào generate(), hàm này sẽ gây ra lỗi nếu mô hình không có các lớp MTP. Giống như speculative decoding, MTP hỗ trợ tìm kiếm tham lam và lấy mẫu nhưng không hỗ trợ đầu vào dạng batch.

Ở đây, MTP chỉ các lớp dự đoán được bao gồm trong checkpoint đích, như trong DeepSeek-V3 và GLM-4.5. Gemma 4 sử dụng một checkpoint trợ lý riêng cho việc giải mã MTP. Truyền nó dưới dạng assistant_model (xem Drafter dự đoán đa token).

Pipeline text-generation và generate() đều chấp nhận use_mtp=True.

from transformers import pipeline

pipe = pipeline(
    "text-generation",
    model="zai-org/GLM-4.5-Air",
    device_map="auto",
)

outputs = pipe(
    "Hugging Face is an open-source company",
    use_mtp=True,
    do_sample=False,
    max_new_tokens=50,
    add_special_tokens=True,
    return_full_text=True,
)

[output["generated_text"] for output in outputs]

Hoặc bạn có thể suy luận như sau.

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("zai-org/GLM-4.5-Air")
model = AutoModelForCausalLM.from_pretrained("zai-org/GLM-4.5-Air", device_map="auto")
inputs = tokenizer("Hugging Face is an open-source company", return_tensors="pt").to(model.device)

outputs = model.generate(**inputs, use_mtp=True, do_sample=False, max_new_tokens=50)
tokenizer.batch_decode(outputs, skip_special_tokens=True)

Bộ tạo dự đoán đa token

MTP drafter là một mô hình trợ lý nhẹ được nạp từ một checkpoint riêng biệt. Nó đề xuất các token ứng viên để mô hình đích xác minh. Drafter sử dụng các embedding đầu vào của mô hình đích và chú ý đến bộ nhớ đệm KV của nó. Nó bắt đầu từ các trạng thái ẩn của mô hình đích, sau đó phác thảo các token theo cách tự hồi quy. Đối với mỗi token tiếp theo, nó đưa trạng thái ẩn đã chiếu của nó trở lại vào drafter. Mô hình đích xác minh các token đã phác thảo trong một lần truyền tiến, giống như trong speculative decoding.

Khác với các lớp MTP, vốn là một phần của checkpoint đích và chia sẻ output head của nó, một MTP drafter là một mô hình trợ lý riêng biệt. Nó sử dụng input embeddings và KV cache của mô hình đích để đề xuất các token.

from transformers import AutoProcessor, AutoModelForCausalLM

processor = AutoProcessor.from_pretrained("google/gemma-4-E4B-it")
target_model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-4-E4B-it",
    dtype="auto",
    device_map="auto",

)

assistant_model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-4-E4B-it-assistant",
    dtype="auto",
    device_map="auto",
)

messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "Write a short joke about saving RAM."},
]

inputs = processor.apply_chat_template(
    messages,
    tokenize=True,
    return_dict=True,
    return_tensors="pt",
    add_generation_prompt=True,
    enable_thinking=False
).to(target_model.device)
input_len = inputs["input_ids"].shape[-1]

outputs = target_model.generate(
    **inputs,
    assistant_model=assistant_model,
    max_new_tokens=256,
)
response = processor.decode(outputs[0][input_len:], skip_special_tokens=False)

processor.parse_response(response)

Giải mã hỗ trợ phổ quát

Universal assisted decoding (UAD) cho phép thực hiện speculative decoding ngay cả khi mô hình chính và mô hình trợ lý có bộ tokenizer khác nhau. Tính năng này cho phép bạn ghép bất kỳ mô hình trợ lý nhỏ nào với mô hình chính. Các token ứng viên được mã hóa lại và thuật toán tính toán dãy con chung dài nhất để đảm bảo phần tiếp theo được căn chỉnh chính xác.

Truyền tokenizer, assistant_tokenizer và assistant_model vào generate() để bật UAD.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

assistant_tokenizer = AutoTokenizer.from_pretrained("double7/vicuna-68m")
tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-9b")
model = AutoModelForCausalLM.from_pretrained("google/gemma-2-9b", dtype="auto")
assistant_model = AutoModelForCausalLM.from_pretrained("double7/vicuna-68m", dtype="auto")
inputs = tokenizer("Hugging Face is an open-source company", return_tensors="pt")

outputs = model.generate(**inputs, assistant_model=assistant_model, tokenizer=tokenizer, assistant_tokenizer=assistant_tokenizer)
tokenizer.batch_decode(outputs, skip_special_tokens=True)
'Hugging Face is an open-source company that is dedicated to creating a better world through technology.'

Xác minh tập hợp tĩnh

Speculative decoding tiêu chuẩn là không mất mát — nó đảm bảo phân phối đầu ra khớp chính xác với mô hình đích. Xác minh tập hợp tĩnh nới lỏng điều này bằng cách xác minh đối với hỗn hợp của phân phối đích và phân phối nháp, làm tăng tỷ lệ chấp nhận (suy luận nhanh hơn) với cái giá là phân phối đầu ra trở thành một hỗn hợp thay vì phân phối đích chính xác. Giá trị khởi đầu được khuyến nghị là 0.7.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM-1.7B")
model = AutoModelForCausalLM.from_pretrained("HuggingFaceTB/SmolLM-1.7B", dtype="auto")
assistant_model = AutoModelForCausalLM.from_pretrained("HuggingFaceTB/SmolLM-135M", dtype="auto")
inputs = tokenizer("Hugging Face is an open-source company", return_tensors="pt")

outputs = model.generate(**inputs, assistant_model=assistant_model, assistant_ensemble_weight=0.7, do_sample=True)
tokenizer.batch_decode(outputs, skip_special_tokens=True)

Điều này yêu cầu các logits ứng viên từ mô hình trợ lý và không được hỗ trợ với giải mã tra cứu prompt.

Tài nguyên

Đọc bài đăng trên blog Assisted Generation: một hướng đi mới hướng tới tạo văn bản có độ trễ thấp để có thêm ngữ cảnh về độ trễ tạo văn bản và assisted generation.

Update on GitHub