Posted inAI

Knowledge Distillation cho LLM: Cách thu nhỏ mô hình để tăng tốc độ và giảm chi phí

AI tutorial - IT technology blog
AI tutorial - IT technology blog

Bối cảnh & Lý do: Thực tế khi vận hành LLM ở quy mô lớn

Sáu tháng trước, đội ngũ của tôi đã gặp phải một nút thắt cổ chai phổ biến nhưng khá nan giải. Chúng tôi đang vận hành một pipeline RAG phức tạp sử dụng GPT-4, và dù chất lượng rất tuyệt vời, nhưng hóa đơn API thì tăng vọt. Quan trọng hơn, độ trễ từ 5 đến 10 giây cho các tác vụ suy luận đơn giản đã làm hỏng trải nghiệm người dùng. Chúng tôi đã thử prompt engineering và fine-tuning cơ bản trên các mô hình nhỏ hơn như Llama-3-8B, nhưng khoảng cách về khả năng suy luận là quá lớn. Các mô hình nhỏ đơn giản là không thể tái hiện được sự tinh tế của các mô hình lớn hơn.

Đó là lúc tôi chuyển hướng sang Knowledge Distillation (KD – Chưng cất tri thức). Thay vì chỉ huấn luyện một mô hình nhỏ trên dữ liệu thô, KD sử dụng một mô hình “Teacher” (Giáo viên – những gã khổng lồ như GPT-4 hoặc Llama-3-70B) để dẫn dắt một mô hình “Student” (Học sinh – như Phi-3 hoặc Llama-3-8B). Mục tiêu là làm cho mô hình học sinh bắt chước logic nội tại của giáo viên, chứ không chỉ là kết quả đầu ra cuối cùng.

Tôi đã áp dụng phương pháp này trong thực tế và kết quả luôn ổn định. Bằng cách chưng cất các khả năng suy luận cụ thể vào một mô hình 8B, chúng tôi đã giảm được 10 lần chi phí inference và cải thiện đáng kể tốc độ token-per-second, trong khi vẫn duy trì được khoảng 94% độ chính xác của mô hình giáo viên trên các tác vụ chuyên biệt. Nếu bạn đang gặp khó khăn trong việc cân bằng giữa trí tuệ của mô hình và ngân sách vận hành, distillation chính là bước đi hợp lý tiếp theo.

Cài đặt: Thiết lập môi trường Distillation

Để thực hiện distillation, bạn cần một môi trường mạnh mẽ có khả năng xử lý hai mô hình cùng lúc (hoặc ít nhất là xử lý đầu ra của giáo viên thành một tập dữ liệu). Tôi thường sử dụng hệ sinh thái Hugging Face vì nó cung cấp các công cụ hoàn thiện nhất cho quy trình này.

Trước tiên, hãy đảm bảo bạn có máy tính với đủ VRAM. Nếu bạn đang thực hiện distillation từ mô hình 70B sang mô hình 8B tại chỗ (locally), bạn có thể sẽ cần GPU A100 kép hoặc ít nhất là sử dụng các phiên bản quantized của mô hình giáo viên. Đây là thiết lập cơ bản mà tôi thường dùng:

pip install -U torch transformers datasets accelerate bitsandbytes peft trl

Chúng tôi sử dụng trl (Transformer Reinforcement Learning) và peft vì việc distillation toàn bộ tham số (full parameter distillation) thường quá tốn kém. Phần lớn thời gian, tôi áp dụng LoRA (Low-Rank Adaptation) cho mô hình học sinh trong quá trình distillation để tiết kiệm bộ nhớ và thời gian.

Cấu hình: Triển khai quy trình Teacher-Student

Có nhiều cách để chưng cất tri thức, nhưng đối với LLM, chúng ta thường tập trung vào “Response-based Distillation” (Distillation dựa trên phản hồi) hoặc “Logit-based Distillation” (Distillation dựa trên logit). Theo kinh nghiệm của tôi, response-based distillation (thường được gọi là SFT trên dữ liệu do giáo viên tạo ra) là phương pháp thực tế nhất cho các đội ngũ triển khai sản phẩm.

1. Tạo tập dữ liệu Distillation

Bước đầu tiên là yêu cầu mô hình Giáo viên tạo ra các lời giải thích chi tiết cho tập huấn luyện của bạn. Nếu bạn muốn học sinh giỏi về logic, giáo viên phải trình bày các bước thực hiện. Điều này thường được gọi là “Chain-of-Thought Distillation”.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# Tải mô hình Giáo viên (Ví dụ: Llama-3-70B-Instruct)
teacher_id = "meta-llama/Meta-Llama-3-70B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(teacher_id)
model = AutoModelForCausalLM.from_pretrained(
    teacher_id, 
    device_map="auto", 
    load_in_4bit=True # Sử dụng 4-bit để phù hợp với phần cứng tiêu dùng hoặc doanh nghiệp tầm trung
)

def generate_teacher_rationale(prompt):
    messages = [
        {"role": "system", "content": "Bạn là một giáo viên chuyên gia. Hãy đưa ra lập luận chi tiết từng bước cho câu hỏi sau đây."},
        {"role": "user", "content": prompt}
    ]
    # Logic khởi tạo tiêu chuẩn ở đây...
    return teacher_output

2. Huấn luyện mô hình Học sinh

Khi đã có tập dữ liệu gồm { "prompt": "...", "teacher_explanation": "...", "final_answer": "..." }, bạn tiến hành huấn luyện mô hình học sinh. Cấu hình dưới đây hướng dẫn cách thiết lập vòng lặp huấn luyện bằng SFTTrainer, phương pháp mà tôi thấy là đáng tin cậy nhất cho việc này.

from trl import SFTTrainer
from transformers import TrainingArguments

# Mô hình Học sinh: Llama-3-8B
student_id = "meta-llama/Meta-Llama-3-8B"

training_args = TrainingArguments(
    output_dir="./distilled-model",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    learning_rate=2e-4,
    num_train_epochs=3,
    logging_steps=10,
    optim="paged_adamw_32bit",
    fp16=True,
)

trainer = SFTTrainer(
    model=student_id,
    train_dataset=distilled_dataset,
    dataset_text_field="teacher_explanation", # Học sinh học theo logic của giáo viên
    max_seq_length=2048,
    args=training_args,
)

trainer.train()

Bằng cách huấn luyện trên teacher_explanation thay vì chỉ final_answer, mô hình học sinh sẽ học được các mô hình tư duy tiềm ẩn. Đây là lý do tại sao một mô hình 8B được chưng cất thường vượt trội hơn so với mô hình 8B tiêu chuẩn được huấn luyện trên dữ liệu thô.

Xác minh & Giám sát: Đo lường thành công

Sau khi quá trình distillation hoàn tất, bạn không thể chỉ dựa vào các chỉ số loss (mất mát) tiêu chuẩn. Loss thấp trong khi huấn luyện không phải lúc nào cũng đồng nghĩa với một mô hình “thông minh” khi triển khai thực tế. Tôi sử dụng hệ thống xác minh ba cấp để đảm bảo mô hình học sinh đã sẵn sàng để deploy.

1. Sử dụng LLM làm giám khảo (LLM-as-a-Judge)

Tôi thường lấy 500 mẫu thử nghiệm và yêu cầu GPT-4o so sánh đầu ra của Giáo viên và Học sinh. Chúng tôi xem xét “tỷ lệ thắng” (win rates). Nếu Học sinh thắng hoặc hòa với Giáo viên hơn 80% số lần trên các tác vụ chuyên biệt, tôi coi đó là một thành công.

2. Đo chuẩn (Benchmarking) độ trễ và thông lượng

Mục đích chính của việc này là hiệu suất. Tôi sử dụng vLLM để kiểm tra thông lượng của mô hình đã được chưng cất. Trong lần chuyển đổi hệ thống gần đây, chúng tôi đã thấy sự thay đổi sau:

  • Giáo viên (70B): 15 tokens/giây, 0,80 USD trên 1 triệu token.
  • Học sinh (8B Distilled): 95 tokens/giây, 0,05 USD trên 1 triệu token.

3. Giám sát độ lệch (Drift Monitoring)

Sau khi triển khai, tôi giám sát “Điểm tin cậy” (Confidence Score) của mô hình học sinh. Nếu học sinh đưa ra câu trả lời với xác suất log (log-probability) thấp so với những gì chúng ta thấy trong quá trình distillation, nó sẽ kích hoạt cơ chế dự phòng chuyển sang mô hình giáo viên. Cách tiếp cận kết hợp (hybrid) này đảm bảo rằng nếu mô hình học sinh gặp phải nội dung nó không biết, “bộ não lớn” (Giáo viên) sẽ tiếp quản.

Distillation không phải là công việc làm một lần là xong. Khi dữ liệu của bạn thay đổi, bạn nên định kỳ chạy lại pipeline distillation để nắm bắt những sắc thái mới từ giáo viên. Đó là một chu kỳ cải tiến liên tục giúp giữ cho chi phí vận hành thấp mà không làm giảm chất lượng mà người dùng mong đợi.

Share:
Tags: