Дистилляция LLM: как сделать модель в 10 раз быстрее
opensourceaillmdistillationoptimizationcompressionit
Введение: Model Distillation
Проблема
Large LLMs are slow and expensive:
- Llama-3-70B: 140 GB, 50 tok/s on A100
- GPT-4: $30/1M input tokens
- Latency: 500ms+ TTFT
Solution: Model Distillation
- Train small model to mimic large model
- 10x speedup, 90% quality retention
- 10x cost reduction
Факт: Distilled models can achieve 90% of teacher quality with 10x inference speed.
Что такое Distillation?
Основная идея
Knowledge Distillation:
Teacher (Large Model):
┌─────────────────────┐
│ 70B parameters │
│ High quality │
│ Slow inference │
│ Expensive │
└─────────────────────┘
↓ labels + logits
Student (Small Model):
┌─────────────────────┐
│ 3B parameters │
│ ~90% quality │
│ Fast inference │
│ Cheap │
└─────────────────────┘
Types of Distillation
1. Knowledge Distillation (KD):
- Student learns from teacher's logits
- Soft labels contain more info than hard labels
- Temperature scaling for softer probabilities
2. Fine-Tuning on Teacher Outputs:
- Generate data with teacher
- Fine-tune student on this data
- Simple but effective
3. Logit Distillation:
- Match student logits to teacher logits
- KL divergence loss
- Works well for classification
4. Feature Distillation:
- Match intermediate representations
- Layer-by-layer knowledge transfer
- More detailed knowledge transfer
Knowledge Distillation
Temperature Scaling
import torch
import torch.nn as nn
import torch.nn.functional as F
def soft_targets(logits, temperature=3.0):
"""Apply temperature to logits for softer probabilities"""
scaled_logits = logits / temperature
return F.softmax(scaled_logits, dim=-1)
def distillation_loss(student_logits, teacher_logits, temperature=3.0):
"""KL divergence between student and teacher"""
student_soft = soft_targets(student_logits, temperature)
teacher_soft = soft_targets(teacher_logits, temperature)
loss = F.kl_div(
F.log_softmax(student_logits / temperature, dim=-1),
teacher_soft,
reduction="batchmean"
) * (temperature ** 2)
return loss
Loss Function
class DistillationLoss(nn.Module):
def __init__(self, temperature=4.0, alpha=0.5):
super().__init__()
self.temperature = temperature
self.alpha = alpha # weight for distillation loss
def forward(self, student_logits, teacher_logits, labels):
# Standard cross-entropy loss
ce_loss = F.cross_entropy(student_logits, labels)
# Distillation loss
distill_loss = distillation_loss(
student_logits, teacher_logits, self.temperature
)
# Combined loss
total_loss = (1 - self.alpha) * ce_loss + self.alpha * distill_loss
return total_loss
# Typical values:
# temperature: 3-8 (higher = softer)
# alpha: 0.5-0.9 (more weight on distillation)
Training Loop
# Training loop for distillation
model_student = AutoModelForCausalLM.from_pretrained("model-3b")
model_teacher = AutoModelForCausalLM.from_pretrained("model-70b")
model_teacher.eval() # Freeze teacher
optimizer = AdamW(model_student.parameters(), lr=5e-5)
criterion = DistillationLoss(temperature=4.0, alpha=0.7)
for batch in dataloader:
# Teacher inference (no grad)
with torch.no_grad():
teacher_outputs = model_teacher(**batch)
# Student inference
student_outputs = model_student(**batch)
# Compute loss
loss = criterion(
student_outputs.logits,
teacher_outputs.logits,
batch["labels"]
)
# Backward pass
loss.backward()
optimizer.step()
optimizer.zero_grad()
Fine-Tuning on Teacher Outputs
Data Generation
# Generate training data with teacher model
def generate_teacher_data(model, prompts, num_samples=10000):
"""Generate responses using teacher model"""
data = []
for prompt in random.sample(prompts, num_samples):
# Generate with teacher
inputs = tokenizer(prompt, return_tensors="pt")
output = model.generate(
**inputs,
max_new_tokens=512,
temperature=0.7,
do_sample=True
)
response = tokenizer.decode(output[0], skip_special_tokens=True)
data.append({
"prompt": prompt,
"response": response,
})
return data
# Save for student training
import json
with open("teacher_data.json", "w") as f:
json.dump(data, f)
Student Training
# Fine-tune student on teacher-generated data
from datasets import Dataset
# Load teacher-generated data
teacher_data = Dataset.from_json("teacher_data.json")
# Tokenize
def tokenize(example):
return tokenizer(
example["prompt"] + example["response"],
truncation=True,
max_length=2048
)
tokenized_data = teacher_data.map(tokenize, batched=True)
# Train student
from transformers import Trainer, TrainingArguments
training_args = TrainingArguments(
output_dir="./student-model",
num_train_epochs=3,
per_device_train_batch_size=8,
learning_rate=2e-5,
fp16=True,
logging_steps=100,
)
trainer = Trainer(
model=model_student,
args=training_args,
train_dataset=tokenized_data,
)
trainer.train()
LLM Distillation in Practice
Open Source Examples
Popular distilled models:
1. Phi-2 (Microsoft):
- Teacher: 13B+ model
- Student: 2.7B parameters
- Trained on "synthetic text"
- Beats GPT-3 on many benchmarks
2. MiniCPM (OpenBMB):
- Teacher: Llama-70B
- Student: 2B parameters
- Matches GPT-3.5 on some tasks
3. Qwen2.5 (Alibaba):
- Distilled variants available
- 0.5B, 1.5B, 3B, 7B sizes
- Strong performance for size
4. Distill-Llama (community):
- Various community efforts
- 70B → 7B distillation
- ~85% quality retention
Quality Comparison
Model Size MMLU HumanEval Speed
────────────────────────────────────────────────────
GPT-4 - 88.4 67.3 Slow
Llama-3-70B 70B 82.0 75.2 1x
Llama-3-8B 8B 73.2 62.1 5x
Distilled-7B 7B 76.5 68.4 6x
Distilled-3B 3B 71.8 58.9 12x
Distilled-1B 1B 65.2 45.3 25x
Speed relative to 70B model
Advanced Techniques
Self-Distillation
Self-distillation (no external teacher):
1. Train initial model
2. Generate data with trained model
3. Fine-tune on generated data
4. Repeat (iterative)
Benefits:
- No external teacher needed
- Improves consistency
- Can generate unlimited data
Risks:
- Model collapse over time
- Need fresh data periodically
Multi-Teacher Distillation
Multiple teachers → one student:
Teachers:
- Llama-3-70B
- Mistral-8x7B
- Mixtral-8x22B
Student:
- 7B model
- Learns from all teachers
- More diverse knowledge
Loss:
loss = Σ α_i * KL(student || teacher_i)
where Σ α_i = 1
Instruction Distillation
Focus on instruction-following:
1. Generate instruction-response pairs:
- Use teacher on instruction dataset
- Collect high-quality responses
2. Fine-tune student:
- Supervised fine-tuning on teacher responses
- Add RLHF for alignment
3. Result:
- Small model with strong instruction following
- Phi-3-mini matches GPT-3.5 on MT-Bench
Compression Beyond Distillation
Combined Approaches
Distillation + Quantization + Pruning:
Original: 70B FP16
140 GB, 100% quality
Step 1: Distillation → 7B
14 GB, 90% quality
Step 2: Quantization INT4 → 7B
3.5 GB, 87% quality
Step 3: Pruning 10% → 6.3B
3.1 GB, 85% quality
Result: 10x smaller, 85% quality, 20x faster
Structured Pruning
# Prune less important weights
import torch.nn.utils.prune as prune
model = AutoModelForCausalLM.from_pretrained("model-7b")
# Prune attention heads
for name, module in model.named_modules():
if "attention" in name and isinstance(module, nn.Linear):
# Prune 10% of weights
prune.l1_unstructured(module, name='weight', amount=0.1)
# Save pruned model
model.save_pretrained("./pruned-model")
tokenizer.save_pretrained("./pruned-model")
Evaluation
How to Measure Success
Distillation metrics:
1. Quality retention:
student_score / teacher_score
2. Speedup:
teacher_latency / student_latency
3. Efficiency score:
quality_retention * speedup
4. Cost reduction:
teacher_cost / student_cost
Target:
- Quality retention > 85%
- Speedup > 5x
- Efficiency score > 1.0
Benchmark Comparison
Benchmark Teacher Student Retention
─────────────────────────────────────────────────
MMLU 82.0 76.5 93.3%
HumanEval 75.2 68.4 91.0%
GSM8K 83.8 78.2 93.3%
MT-Bench 8.2 7.5 91.5%
Заключение
Distillation is powerful for LLM deployment:
Key techniques:
- Knowledge distillation (logits matching)
- Fine-tuning on teacher outputs
- Self-distillation
- Multi-teacher distillation
Results:
- 90% quality with 10x speedup
- Phi-3-mini matches GPT-3.5
- Community distillations available
Best practices:
- Use temperature 3-8 for soft targets
- Combine with quantization for max compression
- Evaluate on task-specific benchmarks
- Monitor for model collapse in self-distillation