Pruning, Knowledge Distillation, Benchmarking

production model 더 작고 빠르게 만드는 세 더 기법

Pruning

Pruning 이 중요하지 않다고 판단된 weight 를 zero (크기 또는 structured rule). 두 맛:

Unstructured — 개별 weight zero. 이론적으로 flops 감소하지만 대부분 hardware 가 sparse matmul 안 가속, 그래서 실용 win 작음.
Structured — 전체 channel / head / block 제거. parameter 더 적음 AND flops 더 적음 AND activation 더 작음. 실용 win 있는 곳.

Pruning + quantization 잘 결합. Pruning + distillation (pruned teacher 에서 더 작은 student train) 더 잘 결합.

Knowledge distillation

큰 'teacher' model 의 soft prediction 일치하게 작은 'student' model train. student 가 hard label 아니라 teacher 의 logit 에서 학습 — 더 많은 정보 운반 ('이건 대부분 cat 이지만 살짝 fox'). 표준 recipe:

teacher 를 input 에 돌리기 — soft logit 받기.
student 를 input 에 돌리기 — 자기 logit 받기.
Loss = α · KL(student / T || teacher / T) · T² + (1 − α) · CE(student, hard_label)

여기서 T 가 'temperature' (보통 2-8) 가 distribution 둘 다 부드럽게, α 가 soft vs hard target weight.

Benchmarking — 너 정직한 성적표

모든 최적화 주장이 benchmark 받을 자격. 실제 deploy batch size 에, 실제 hardware 에, 실제 input distribution 으로 latency 측정. 일화적 speedup 이 production 접촉에 안 살아남음.

Code

Unstructured pruning — magnitude-based·python

import torch
import torch.nn as nn
import torch.nn.utils.prune as prune

linear = nn.Linear(100, 50)

# Zero out 30% of weights by magnitude (smallest go to zero)
prune.l1_unstructured(linear, name='weight', amount=0.3)

# Check sparsity
sparsity = (linear.weight == 0).float().mean().item()
print(f"sparsity: {sparsity:.2%}")           # ~30%

# Make pruning permanent (removes the mask, keeps the zeros)
prune.remove(linear, 'weight')

Knowledge distillation — 고전 recipe·python

import torch
import torch.nn as nn
import torch.nn.functional as F

teacher = BigModel().eval()                   # frozen
student = SmallModel().train()
optimizer = torch.optim.AdamW(student.parameters(), lr=1e-3)
T = 4.0                                        # temperature
alpha = 0.7                                    # soft-vs-hard weight

for x, y in train_loader:
    optimizer.zero_grad()

    with torch.inference_mode():
        teacher_logits = teacher(x)

    student_logits = student(x)

    soft = F.kl_div(
        F.log_softmax(student_logits / T, dim=-1),
        F.softmax(teacher_logits / T, dim=-1),
        reduction='batchmean',
    ) * (T ** 2)

    hard = F.cross_entropy(student_logits, y)

    loss = alpha * soft + (1 - alpha) * hard
    loss.backward()
    optimizer.step()

Benchmark utility — 모두 필요한 표·python

import time
import torch

@torch.inference_mode()
def benchmark(model, input_tensor, num_runs=100, warmup=10, label='model'):
    is_cuda = next(model.parameters()).is_cuda
    if is_cuda: torch.cuda.synchronize()

    for _ in range(warmup):
        model(input_tensor)
    if is_cuda: torch.cuda.synchronize()

    t0 = time.perf_counter()
    for _ in range(num_runs):
        model(input_tensor)
    if is_cuda: torch.cuda.synchronize()
    elapsed = time.perf_counter() - t0

    avg_ms = elapsed / num_runs * 1000
    throughput = input_tensor.size(0) * num_runs / elapsed
    print(f"{label:20s} avg {avg_ms:6.2f} ms/call, {throughput:,.0f} samples/sec")
    return avg_ms

# Use it like:
benchmark(fp32_model, batch, label='fp32 baseline')
benchmark(quant_model, batch, label='int8 quant')
benchmark(compiled_quant_model, batch, label='int8 + compile')

Pruning, Knowledge Distillation, Benchmarking

production model 더 작고 빠르게 만드는 세 더 기법

Pruning

Knowledge distillation

Benchmarking — 너 정직한 성적표

Code

External links

Exercise

Progress

댓글 0