AI 경량화 - 더 빠르고 저렴한 AI 서비스를 위해(NAVER 강의)
1. 서론
AI 경량화는 사실 세간의 인식에 비해 쉽다
기저에 깔려있는 이론은 어렵지만 적용하기에는 매우 쉽다
현재 AI모델은 더 큰 모델, 더 큰 파라미터로 더 좋은 성능을 내는 것이 트렌드
NLP 뿐만이 아니라 CV도 마찬가지
하지만 문제는 서비스 응답 목표치에 비해 AI모델의 추론 속도가 매우 느리다는거
경량화를 통해 AI모델의 아키텍처는 그대로, 정확도 손실은 거의 없게
그런데 추론 속도를 4배 더 빠르게 할수 있다면?
Clova의 LaRva 모델은 1배치당 평균 175.87ms인데 경량화를 통해 4배 더 빠른 43.86ms를 달성했다고함
이 정도면 서빙하고도 남는 수준
경량화 기법의 계통
pruning과 low rank 기법이 주로 연구되고 있고(2023.10 기준)
knowledge distillation, quantization 기법은 뜸하긴 하나 가끔 효율적인 경량화 기법이 소개되는
이 외에 HW target optimization, AI kernel optimization,
remodeling(완성된 모델의 일부 아키텍처를 바꿔 정확도 손실은 거의 없고 속도 이점을 챙기는 기법)
2. quantization
가장 먼저 고려해볼만한 기법
적용도 쉽고 리소스도 적게 들어가며 효과는 가장 좋은
dynamic quantization 예시
quantized_model = torch.quantization.quantize_dynamic(
original_model,
{torch.nn.Linear},
dtype = torch.qint8
)
float32 파라미터를 int8로 바꾸기만 해도 모델 크기는 4배 줄고 속도는 2~3배 빨라짐
완성된 모델의 자료형만 바꾸므로 프레임워크에 종속적이지 않음
pruning과도 궁합이 좋고 서로가 정확도에 악영향을 끼치지 않는다는 논문 실험결과가 많다
적용하면서 겪을 수 있는 이슈는 보통 2가지
quantization이 적용되지 않는 레이어가 존재해서, 이것때문에 발생하는 오류
만약 알수없는 오류로 quantization이 실패했다면, 지원되지 않는 레이어인지 확인해봐야함
본인이 직접 추가하는 것을 검토함하는게 좋을수도
두번째는 torch.jit.script 모듈로 감싸져있는 경우 quantization이 실패할 수 있다
이는 jit.script 모듈로 변환하기 전에 quantization을 실시하고 torch.jit.script 모듈로 변환해야함
jit.script 모듈로 변환되기 전 statedict 파일을 가지고 있어야겠지
3. pruning
학계에서 가장 활발히 연구하고 있는 포텐셜이 가장 높은 분야
하지만 실제 적용하기는 어렵다
완성된 모델의 레이어 하나하나에 target sparsity에 맞춰 prune 함수를 적용하면 끝
import torch
import torch.nn.utils.prune as prune
def prune_bert(model,sparsity = 0.9):
#embeddings
prune.l1_unstructured(model.embeddings.word_embeddings, name = 'weight', amount = sparsity)
prune.l1_unstructured(model.embeddings.position_embeddings, name = 'weight', amount = sparsity)
prune.l1_unstructured(model.embeddings.token_type_embeddings, name = 'weight', amount = sparsity)
#encoder layer
for layer in model.encoder.layer:
#attention weights
prune.l1_unstructured(layer.attention.self.query, name = 'weight', amount = sparsity)
prune.l1_unstructured(layer.attention.self.key, name = 'weight', amount = sparsity)
prune.l1_unstructured(layer.attention.self.value, name = 'weight', amount = sparsity)
prune.l1_unstructured(layer.attention.output.dense, name = 'weight', amount = sparsity)
#feed forward weights
prune.l1_unstructured(layer.intermediate.dense, name = 'weight', amount = sparsity)
prune.l1_unstructured(layer.output.dense, name = 'weight', amount = sparsity)
prune.l1_unstructured(model.pooler.dense, name = 'weight', amount = sparsity)
return model
pruning을 진행하면, 모델의 구조가 취약해질 수 있다고 생각하기 쉽다
즉 pruning을 하면 파라미터 수가 줄어들 것으로 예상
하지만 실제로 pruning을 하면 파라미터 수는 동일한데??? 진짜로?
왜냐하면 이론과는 다르게 실제로 파라미터를 잘라내는 것이 아니라
원래 가지고 있던 값을 0으로 만드는 zero화 시키는 방법을 적용
모델의 내부 텐서 변화를 보면
0이 차지하는 비율을 sparsity라고 하고 반대 개념은 dense로 0이 아닌 값이 차지하는 정도
따라서 tensor의 전체 구조 shape에는 변화가 없다
바뀐 것이 없고 0으로만 되었을 뿐인데 계산속도가 왜 빨라지는가?
sparse matrix, 즉 0이 많은 matrix에는 다른 계산 알고리즘이 적용된다
cpu, gpu에 따라 30%정도 속도 차이가 난다는데
근데 실제로는 오히려 더 느려짐
이는 아주 흔한 문제이다
이럴 떄는 masking된 original weight가 있는지 확인해봐야함
torch, tensorflow에서 prune시 원래 weight와 prune weight를 동시에 가지고 있음
torch.nn.utils.prune.remove함수로 original weight를 제거해야 정상적으로 prune된 weight만 남는다
그래도 문제는 해결이 안된다던데
pytorch, tensorflow에서 sparsity accleration이 지원되지 않기 때문이라함
intel cpu같은 경우 intel mkl같은 sparsity accleration을 지원하는 수학 라이브러리를 설치해야한다함
이런 이슈들을 해결하면 6배~12배 빨라지는 경우를 볼 수 있다함
그럼에도 불구하고 이미 train된 네트워크의 weight를 잘라내는 것이 정확도에 악영향을 끼칠 것 같다는 건데
실제로도 naive하게 pruning을 하면 주요 지표에 악영향을 끼치는건 사실
그러면 왜 네트워크는 일반적으로 over parameterized되어서 pruning을 하더라도 정확도에 영향이 없다고 하는 것인가
pruning의 연구 방향은 쓸모 없는 가중치를 잘라내고 중요한 가중치만 남겨놓아 모델의 성능은 유지하고 속도를 높이는 방식
실제 효과가 있었다는 방식은 2019 lottery ticket hypothesis 논문에 나온 중요 가중치를 식별하여 보존하는 방법
2015 learning both weights and connections for efficient neural network에 나온 iterative pruning
좀 더 좋은건 2015 iterative pruning이라는데
이미 완성된 모델을 target sparsity로 pruning하고 3epoch정도 retrain 시키는 것을 몇번 반복
위 그래프를 보면 pruning을 했는데 오히려 정확도가 높아지는 경우가 있다
불필요한 가중치를 제거하고 중요한 가중치만 남겨 retrain 시키니 정확도가 높아진다는 설명
이건 인간도 마찬가지
어렸을때 비해 어른일때 시냅스가 반토막 나지만 성인일때 오히려 전문적임
실제로 pruning된 네트워크는 cpu에서 3배정도 빠르고 전력도 7배나 덜 소모
4. low rank methodology
혁신적이고 놀라운 방법들이 많지만 적용이 까다롭고 아직 연구가 덜 된 분야
가장 먼저 시도해볼 부분은 low rank approximation을 이용한 행렬 근사
def optimal_rank(tensor, energy_threshold = 0.95):
if len(tensor.shape) != 2:
raise ValueError("Input tensor should be 2D")
_, s, _ = torch.svd(tensor)
total_energy = torch.sum(s**2)
cumulative_energy = torch.cumsum(s**2, dim = 0)/total_energy
optimal_rank = torch.searchsorted(cumulative_energy, energy_threshold) + 1
return int(optimal_rank.item())
def low_rank_approximation(model):
for name, param in model.named_parameters():
if len(param.shape) == 2:
rank = optimal_rank(param, energy_threshold = 0.95)
u,s,v = torch.svd(param)
u_approx = u[:,:rank]
s_approx = s[:rank]
v_approx = v[:,:rank]
param_approx = u_approx @ torch.diag(s_approx) @ v_approx.T
param.data = param_approx
return model
singular value decomposition을 이용해서 bottleneck convolution의 속도를 2~3배 높이고
layer의 메모리 사용량을 5~13배 줄였다는 논문 결과가 있다
하지만 eigenvalue decomposition이라는 행렬 분해 기법인 만큼 결과를 확신할 만큼 선례가 많지 않았고
최적화에 어려움을 줄 수 있어 stable service에 적용하기 어려워보여
하지만 최근 LoRA나 QLoRA 등 신기법이 LLM과 함께 적용되면서 아주 적은 양의 parameter의 retrain만으로
base 모델의 output을 개인화 시킬 수 있다는 점에서 혁신적인 포인트가 많다
5. knowledge distillation
생각과는 다르게 써보세요?
더 크고 더 정밀한 teacher model에서 학습한 가중치를 더 작은 student model에 전이 시키는 것
transfer learning과 비슷하지만, transfer learning은 서로 다른 도메인에서 지식을 전달하는 방식으로 accuracy improvement
knowledge distillation은 서로 같은 도메인에서 모델 A가 모델 B에게 지식을 전달하므로 model compression
import torch
from transformers import BertTokenizer, BertForSequenceClassification, AdamW
from torch.nn import KLDivLoss, Softmax, LogSoftmax
teacher_model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
student_model = BertForSequenceClassification.from_pretrained('bert-base-uncased',num_labels = 2)
teacher_model.eval()
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
temperature = 2.0
criterion = KLDivLoss(reduction = 'batchmean')
softmax = Softmax(dim = -1)
logsoftmax = LogSoftmax(dim = -1)
optimizer = AdamW(student_model.parameters(), lr = 1e-5)
for epoch in range(3):
for batch in dataloader:
inputs = batch['input_ids']
attention_mask = batch['attention_mask']
with torch.no_grad():
teacher_outputs = teacher_model(input_ids = inputs, attention_mask = attention_mask)
teacher_probs = softmax(teacher_outputs.logits/temperature)
student_outputs = student_model(input_ids = inputs, attention_mask = attention_mask)
student_probs = logsoftmax(student_outputs.logits/temperature)
loss = criterion(student_probs, teacher_probs.detach())
loss.backward()
optimizer.step()
optimizer.zero_grad()
기존의 상식은 "더 크고 정밀한 모델의 가중치 전이"였지만
최신 트렌드는 완성된 모델의 정확도 보간법으로 이용되고 있다
self distillation은 학습, 경량화가 모두 적용된 완성된 모델을 retrain하여 정확도를 1%라도 높이는 기법
실제 라이브 서비스 모델에 적용했을때 0.3% 정확도를 높였다고 함
이는 hyperparameter 조정이나 모델의 구조적 변경보다 리소스적으로 경제적이면서도 확실함
반대로 말하면 효과가 크지 않다
경험적으로 teacher model에서 student model로의 knowledge distillation은 효과가 거의 없거나 역효과
제대로 성능을 볼려면 많은 리소스가 들어감
차라리 완성된 모델의 self distillation을 통해 마지막 모델의 성능 쥐어짜는게 제일 좋다
6. benefit
AI도 CPU, GPU, TPU등 특정 프로세서에서 전기적 신호에 의해 이루어지는 연산 과정으로, 전력 소모, 비용을 생각해야함
경량화는 전력 소모 비용 절감에 도움을 줌
AI모델의 추론 속도를 빠르게 하고 하드웨어 장비에 적은 부하
하지만 Naive한 방법의 경량화는 성능을 낮출 수 있음
이런 경우 충분히 연구된 경량화 기법을 적용하여 성능을 오히려 향상시킬 수도 있음
7. risk
1) 일반화의 어려움
AI 모델은 input이 네트워크에 들어가 output이 나온다는 공통점이 있지만
데이터셋, 모델 아키텍처에 따라 내부적인 차이가 존재하고 이로 인해 특정 모델에는 효과적인 경량화 기법이
다른 모델에서는 아무런 효과가 없거나 아니면 오히려 역효과
pruning > 학습된 모델이 특정 뉴런에 중요 정보가 몰려있는 경우 해당 뉴런을 자르면 속도적인 이점도 못얻고 성능도 떨어짐
quantization > 네트워크에 정보가 전반에 큰 크기로 저장된 경우 데이터 타입을 강제로 줄이는 것은 피처를 상당히 손실
knowledge distillation > teacher, student 모델의 구조가 크게 다르면 weight가 제대로 전이되지 않을 수 있음
2) 아직 발전하지 않은 분야
논문도 적고, 지식 습득도 어렵고, 노하우에 의존되며 신규 논문 지식을 빠르게 습득해야
정형화된 방법이 없어 try and error로 시도해봐야
그럼에도 불구하고... 얻을 수 있는 이점은?
1) 성능, 속도 변화를 설명가능하고 예측가능함
특정 경량화 기법이 예상대로 동작하지 않으면 기술적으로 해석 가능하고 개선 포인트가 명확함
불확실성이 감소하고 기술적으로 stable
2) 직관적이고 안정적인 output
결과가 서비스에 가져올 영향을 누구나 이해 가능함
accuracy나 latency가 얼마나 변화하는가?
cpu나 메모리 점유율이 얼마나 줄어드는지
검증된 모델 아키텍처를 최대한 바꾸지 않으면서 속도와 성능을 높이니 안정적
3) 확장성 지속성
이미 서비스 가능한 모델일지라도 경량화 기법을 적용하여 더 나은 서비스를 제공할 수 있다는 점
https://www.youtube.com/watch?v=NVNCPGWe5Ss