Дистилляция LLM: Сжатие больших моделей в маленькие
llmdistillationcompressionmodel-sizingoptimization
Введение
Дистилляция знаний (Knowledge Distillation) — это техника сжатия больших языковых моделей (учителей) в маленькие модели (студенты) с минимальной потерей качества. Это один из самых эффективных способов адаптации LLM для production-использования.
Теория дистилляции
Основная идея
┌─────────────────────────────────────────────────────────┐
│ Training Phase │
├─────────────────────────────────────────────────────────┤
│ │
│ Large Model (Teacher) Small Model (Student) │
│ ┌──────────────┐ ┌──────────────┐ │
│ │ 70B / 175B │ │ 7B / 1.5B │ │
│ │ FP16 │ │ INT8/INT4 │ │
│ └──────┬───────┘ └──────┬───────┘ │
│ │ Soft Labels │ Soft Labels │
│ │ (logits) │ (learned) │
│ └──────────┬───────────────┘ │
│ ↓ │
│ Distillation Loss │
│ L = α·L_task + (1-α)·L_distill │
│ │
├─────────────────────────────────────────────────────────┤
│ Inference Phase │
├─────────────────────────────────────────────────────────┤
│ │
│ Student Model Only │
│ ┌──────────────┐ │
│ │ 7B / 1.5B │ → 10x faster, 10x smaller │
│ │ INT8/INT4 │ │
│ └──────────────┘ │
│ │
└─────────────────────────────────────────────────────────┘
Математическая основа
import torch
import torch.nn as nn
import torch.nn.functional as F
class DistillationLoss(nn.Module):
"""Комбинированная функция потерь для дистилляции"""
def __init__(self, temperature: float = 4.0, alpha: float = 0.7):
super().__init__()
self.temperature = temperature
self.alpha = alpha
self.ce_loss = nn.CrossEntropyLoss()
self.kl_loss = nn.KLDivLoss(reduction="batchmean")
def forward(self, student_logits, teacher_logits, targets):
# Task loss (cross-entropy on hard labels)
hard_loss = self.ce_loss(
student_logits / self.temperature,
targets
)
# Distillation loss (KL divergence on soft labels)
soft_loss = self.kl_loss(
F.log_softmax(student_logits / self.temperature, dim=-1),
F.softmax(teacher_logits / self.temperature, dim=-1)
) * (self.temperature ** 2)
# Combined loss
return self.alpha * hard_loss + (1 - self.alpha) * soft_loss
Типы дистилляции
1. Logit-based Distillation
class LogitDistillation:
"""Дистилляция на уровне логитов"""
def __init__(self, teacher, student, temperature=4.0):
self.teacher = teacher
self.student = student
self.temperature = temperature
self.optimizer = torch.optim.Adam(student.parameters(), lr=1e-4)
def distill_step(self, batch_x, batch_y):
# Teacher inference (frozen)
with torch.no_grad():
teacher_logits = self.teacher(batch_x)
# Student inference
student_logits = self.student(batch_x)
# Compute losses
ce_loss = F.cross_entropy(student_logits, batch_y)
kl_loss = F.kl_div(
F.log_softmax(student_logits / self.temperature, dim=-1),
F.softmax(teacher_logits / self.temperature, dim=-1),
reduction='batchmean'
) * (self.temperature ** 2)
loss = 0.3 * ce_loss + 0.7 * kl_loss
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
return loss.item(), ce_loss.item(), kl_loss.item()
2. Feature-based Distillation
class FeatureDistillation:
"""Дистилляция на уровне скрытых представлений"""
def __init__(self, teacher, student, num_layers=4):
self.teacher = teacher
self.student = student
# Projection layers to match dimensions
self.projections = nn.ModuleList([
nn.Linear(teacher_hidden_dim, student_hidden_dim)
for _ in range(num_layers)
])
self.mse_loss = nn.MSELoss()
def get_intermediate_layers(self, model, x):
"""Получение промежуточных представлений"""
features = []
hidden = x
for name, module in model.named_modules():
if isinstance(module, torch.nn.TransformerEncoderLayer):
hidden = module(hidden)
features.append(hidden)
if len(features) >= 4:
break
return features
def forward(self, x, y):
# Teacher features (frozen)
with torch.no_grad():
teacher_features = self.get_intermediate_layers(self.teacher, x)
# Student features
student_features = self.get_intermediate_layers(self.student, x)
# Feature matching loss
feature_loss = 0
for t_feat, s_feat, proj in zip(teacher_features, student_features, self.projections):
projected_s = proj(s_feat)
feature_loss += self.mse_loss(projected_s, t_feat)
# Task loss
task_loss = F.cross_entropy(self.student(x), y)
return task_loss + 0.5 * feature_loss
3. Response-based Distillation
class ResponseDistillation:
"""Дистилляция на уровне ответов (без учителя)"""
def __init__(self, student, tokenizer):
self.student = student
self.tokenizer = tokenizer
def self_distill(self, texts, confidence_threshold=0.9):
"""Self-distillation: сильный студент → слабый студент"""
# Strong student generates pseudo-labels
strong_outputs = self.student.generate(
texts, max_length=512, do_sample=False
)
# Filter high-confidence predictions
confident_mask = self._check_confidence(strong_outputs)
# Train weak student on filtered outputs
weak_loss = self._train_weak_student(
texts, strong_outputs, confident_mask
)
return weak_loss
4. Imitation Learning
class ImitationDistillation:
"""Дистилляция через имитационное обучение"""
def __init__(self, teacher, student, tokenizer):
self.teacher = teacher
self.student = student
self.tokenizer = tokenizer
def generate_expert_trajectories(self, prompts: List[str], max_len=256):
"""Генерация траекторий эксперта"""
trajectories = []
for prompt in prompts:
input_ids = self.tokenizer.encode(prompt, return_tensors="pt")
with torch.no_grad():
output = self.teacher.generate(
input_ids,
max_length=len(input_ids[0]) + max_len,
do_sample=True,
temperature=0.7
)
trajectories.append(output)
return trajectories
def fine_tune_on_expert(self, trajectories: List[torch.Tensor]):
"""Fine-tuning студента на траекториях эксперта"""
optimizer = torch.optim.AdamW(self.student.parameters(), lr=5e-5)
for trajectory in trajectories:
input_ids = trajectory[:, :-1]
targets = trajectory[:, 1:]
outputs = self.student(input_ids, labels=targets)
loss = outputs.loss
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(self.student.parameters(), 1.0)
optimizer.step()
Advanced Techniques
1. Multi-Teacher Distillation
class MultiTeacherDistillation:
"""Дистилляция от нескольких учителей"""
def __init__(self, teachers: List[nn.Module], student: nn.Module):
self.teachers = teachers
self.student = student
self.teacher_weights = nn.Parameter(
torch.ones(len(teachers)) / len(teachers)
)
def forward(self, x, y):
# Ensemble teacher predictions
teacher_logits = []
for teacher in self.teachers:
with torch.no_grad():
logits = teacher(x)
teacher_logits.append(logits)
# Weighted ensemble
weights = F.softmax(self.teacher_weights, dim=0)
ensemble_logits = sum(
w * logit for w, logit in zip(weights, teacher_logits)
)
# Distill to student
student_logits = self.student(x)
distill_loss = F.kl_div(
F.log_softmax(student_logits / T, dim=-1),
F.softmax(ensemble_logits / T, dim=-1),
reduction='batchmean'
) * (T ** 2)
task_loss = F.cross_entropy(student_logits, y)
return task_loss + 0.7 * distill_loss
2. Curriculum Distillation
class CurriculumDistillation:
"""Дистилляция с учебной программой"""
def __init__(self, teacher, student):
self.teacher = teacher
self.student = student
self.curriculum = self._build_curriculum()
def _build_curriculum(self):
"""Построение учебной программы от простого к сложному"""
return [
{"difficulty": 0.1, "epochs": 5, "data_subset": "simple"},
{"difficulty": 0.3, "epochs": 10, "data_subset": "intermediate"},
{"difficulty": 0.6, "epochs": 15, "data_subset": "complex"},
{"difficulty": 1.0, "epochs": 20, "data_subset": "all"},
]
def train(self, data_loader_by_difficulty):
for stage in self.curriculum:
loader = data_loader_by_difficulty[stage["data_subset"]]
for epoch in range(stage["epochs"]):
for x, y in loader:
self.distill_step(x, y, stage["difficulty"])
3. Layer-wise Distillation
class LayerwiseDistillation:
"""Послойная дистилляция"""
def __init__(self, teacher, student):
self.teacher = teacher
self.student = student
# Freeze student initially
for param in self.student.parameters():
param.requires_grad = False
def distill_layer_by_layer(self):
"""Послойная дистилляция: от низких слоёв к высоким"""
teacher_layers = self._get_teacher_layers()
student_layers = self._get_student_layers()
for i, (t_layer, s_layer) in enumerate(zip(teacher_layers, student_layers)):
# Unfreeze current student layer
for param in s_layer.parameters():
param.requires_grad = True
# Distill this layer
self._distill_layer_pair(t_layer, s_layer)
# Freeze this layer, move to next
for param in s_layer.parameters():
param.requires_grad = False
def _distill_layer_pair(self, teacher_layer, student_layer):
"""Дистилляция пары слоёв"""
# Match output distributions
pass
Practical Examples
Distilling Llama 3 8B → Llama 3 1B
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
# Load teacher
teacher = AutoModelForCausalLM.from_pretrained(
"meta-llama/Meta-Llama-3-8B",
torch_dtype=torch.float16,
device_map="auto"
)
teacher.eval()
# Load student
student = AutoModelForCausalLM.from_pretrained(
"meta-llama/Meta-Llama-3-1B",
torch_dtype=torch.float16,
device_map="auto"
)
# Distillation training
from trl import DPOTrainer, DPOConfig
training_args = DPOConfig(
output_dir="./llama3-8b-distilled",
per_device_train_batch_size=4,
gradient_accumulation_steps=8,
learning_rate=1e-5,
num_train_epochs=3,
logging_steps=10,
save_steps=100,
temperature=4.0, # Temperature for soft labels
)
trainer = DPOTrainer(
model=student,
ref_model=teacher,
args=training_args,
train_dataset=distillation_dataset,
)
trainer.train()
Distillation with vLLM (fast generation)
from vllm import LLM, SamplingParams
# Fast teacher inference with vLLM
teacher_llm = LLM(model="meta-llama/Meta-Llama-3-70B", tensor_parallel_size=4)
sampling_params = SamplingParams(
temperature=0.7,
max_tokens=256,
top_p=0.95,
)
prompts = ["Explain quantum computing:", "Write a poem about AI:"]
teacher_outputs = teacher_llm.generate(prompts, sampling_params)
# Use outputs for student training
for prompt, output in zip(prompts, teacher_outputs):
student_training_data.append({
"prompt": prompt,
"completion": output.outputs[0].text,
})
Performance Comparison
| Teacher | Student | Size Reduction | Quality Retention | Speedup |
|---|---|---|---|---|
| Llama 3 70B | Llama 3 8B | 8.75x | ~92% | ~7x |
| Llama 3 8B | Llama 3 1B | 8x | ~85% | ~6x |
| GPT-3 175B | Distill 3B | 58x | ~78% | ~50x |
| BERT-base | BERT-tiny | 12x | ~90% | ~10x |
Best Practices
- Начинайте с близкого по архитектуре учителя — одинаковая архитектура упрощает дистилляцию
- Используйте несколько стратегий одновременно — logit + feature + response
- Тщательно подбирайте temperature — высокое T = больше информации, низкое T = быстрее обучение
- Алгоритмический отбор данных — дистиллируйте сначала на сложных примерах
- Итеративная дистилляция — цепочка моделей: 70B → 13B → 7B → 1B
Заключение
Дистилляция LLM — это мощный инструмент для создания production-ready моделей. Правильно проведённая дистилляция позволяет сохранить 85-95% качества при 5-10x сокращении размера и ускорении inference.