Pruning нейросетей: обрезка весов для оптимизации и ускорения инференса

Представьте: ваша обученная GPT-подобная модель занимает 16 ГБ и выдаёт ответ за 5 секунд на GPU, а на CPU или edge-устройстве — вообще неприемлемо. Клиенты жалуются на задержки, затраты на инфраструктуру растут. Мы с такой ситуацией сталкиваемся постоянно: модель хороша, но в production её не поставишь — дорого и медленно. Pruning — один из ключевых инструментов, чтобы это исправить.

Pruning — удаление малозначимых параметров (весов, нейронов, attention heads, слоёв) из обученной нейросети. Цель — снизить размер модели и ускорить инференс при минимальной потере качества. В контексте LLM pruning часто комбинируют с квантизацией и дистилляцией для максимального сжатия.

Почему pruning важен для LLM?

Большие языковые модели (LLM) содержат миллиарды параметров, но многие из них избыточны. Pruning позволяет ужать модель в 2–4 раза без значительной деградации, что критично для деплоя на edge, снижения затрат на GPU и уменьшения задержек. Без pruning современные LLM часто остаются исследовательскими артефактами — слишком дорогими для реального применения.

Виды pruning

• Unstructured pruning: обнуляются отдельные веса по всей матрице. Высокое сжатие, но требует sparse computation — стандартные GPU не ускоряют sparse операции «из коробки».
• Structured pruning: удаляются целые структурные элементы — нейроны, attention heads, слои. Результат — реально меньшая плотная модель, которая работает быстрее на стандартном железе.
• Semi-structured pruning (N:M sparsity): удаляются N весов из каждого блока M. Формат 2:4 поддерживается NVIDIA Ampere и выше на аппаратном уровне (до 2× ускорение).

Какой метод pruning выбрать для вашей модели?

Выбор метода pruning зависит от целевого железа и требований к качеству. Ниже — сравнение популярных подходов.

Методы pruning с примерами кода

LLM-Pruner: структурированный pruning LLM

# Пример использования LLM-Pruner # pip install llm-pruner from LLMPruner.pruner import LlamaStructuredPruner from transformers import AutoModelForCausalLM, AutoTokenizer model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3.1-7B") tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3.1-7B") pruner = LlamaStructuredPruner( model=model, tokenizer=tokenizer, pruning_ratio=0.25, # Удалить 25% параметров ) # Вычисление важности параметров на calibration data calibration_data = ["Текст для анализа важности весов...", ...] pruner.get_mask(calibration_data, method="taylor") # Taylor expansion importance # Применение маски и pruning pruned_model = pruner.prune() 

SparseGPT: эффективный unstructured pruning без retraining

Из документации SparseGPT: "SparseGPT can prune 50% of the weights without retraining." — GitHub

# sparsegpt — библиотека от авторов метода # Пример концептуального кода from sparsegpt import SparseGPT sparsegpt = SparseGPT(model) sparsegpt.fasterprune( sparsity=0.5, # 50% sparsity prunen=2, # N в N:M prunem=4, # M в N:M (2:4 — поддерживается аппаратно) percdamp=0.01, blocksize=128, ) 

При 2:4 sparsity (50%) на NVIDIA A100/H100 ускорение inference на Tensor Core около 1.7–2×.

Wanda: простой и эффективный pruning

# Wanda проще SparseGPT, но сопоставимо по качеству # Работает за несколько минут на 7B модели def wanda_pruning(model, calibration_loader, sparsity=0.5): """Упрощённая реализация Wanda""" for name, module in model.named_modules(): if isinstance(module, torch.nn.Linear): # Накапливаем статистику активаций activation_norms = get_activation_norms(module, calibration_loader) # Importance score = |W| * ||X|| importance = module.weight.abs() * activation_norms # Pruning по threshold threshold = torch.quantile(importance, sparsity) mask = importance > threshold module.weight.data *= mask return model 

Depth pruning: удаление слоёв

Для LLM средние слои часто менее критичны, чем первые и последние:

def depth_prune_llm(model, layers_to_remove: list[int]): """Удаление указанных decoder layers""" # Для Llama-архитектуры remaining_layers = [ layer for i, layer in enumerate(model.model.layers) if i not in layers_to_remove ] model.model.layers = torch.nn.ModuleList(remaining_layers) return model # Пример: удаляем 8 средних слоёв из 32 (25% depth reduction) pruned_model = depth_prune_llm(model, layers_to_remove=list(range(12, 20))) # Результат: 24-слойная модель из 32-слойной 

Практический кейс: оптимизация edge-деплоя

Задача: дообученная Llama 3.1 8B для промышленного контроллера (ARM-сервер, 16GB RAM, нет GPU). Требование: инференс < 2с на запрос.

Наша стратегия оптимизации:

1. GGUF Q4_K_M квантизация: 8B → 4.1GB, 8 tok/s на CPU (недостаточно)
2. Depth pruning (удаление 8 слоёв из 32): -25% latency, -3% качества
3. Width pruning attention heads (удаление 20% голов): -15% latency
4. Повторная квантизация: GGUF Q4_K_M на pruned модели

Итоговые характеристики pruned+quantized модели:

В типовом проекте с моделью 7B параметров экономия на инфраструктуре достигает 40%.

# После pruning — краткий fine-tuning для восстановления from trl import SFTTrainer, SFTConfig # Используем тот же датасет, что для fine-tuning, но с меньшим LR recovery_config = SFTConfig( num_train_epochs=1, # 1 эпоха для recovery learning_rate=5e-5, # Ниже, чем при full fine-tuning gradient_checkpointing=True, bf16=True, ) trainer = SFTTrainer(model=pruned_model, args=recovery_config, train_dataset=dataset) trainer.train()