Оптимізація нейромереж: прунінг для прискорення інференсу

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

Прунінг нейромережі — видалення малозначущих параметрів (ваг, нейронів, attention heads, шарів) із навченої нейромережі. Мета — знизити розмір моделі та прискорити інференс при мінімальній втраті якості. У контексті LLM pruning часто комбінують з квантизацією та дистиляцією для максимального стиснення.

Чому pruning важливий для LLM?

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

Види pruning

• Unstructured pruning: обнулюються окремі ваги по всій матриці. Високе стиснення, але потребує sparse computation — стандартні GPU не прискорюють sparse операції «з коробки». Використання тензорних ядер з підтримкою N:M sparsity (2:4) на GPU Ampere+ забезпечує апаратне прискорення до 2×.
• 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%. Наприклад, для моделі Llama 7B на AWS p3.2xlarge ($1.5/год) цілодобова робота коштує ~$1,080/міс, а після pruning — ~$648/міс, що заощаджує $432 щомісяця.

# Після 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()