Shortened LLaMA: A Simple Depth Pruning for Large Language Models

1. Background

최근 연구 동향

• LLM의 발전으로 쳇봇 등 다양한 응용 프로그램이 발전하였음
• 하지만 높은 계산량으로 배포에는 한계가 있음
• 비용 문제 해결을 위해 다양한 연구가 진행되고 있음

Pruning

• LLM 경량화 방법으로 Structed Pruning 이 많이 연구되는 추세
• Pruning에는 Structed Pruning과 Unstructed Pruning, Semi Structed Pruning이 존재
• Structed Pruning에는 Width Pruning, Depth Pruning 존재
• Depth Pruning이 Width Pruning보다 성능이 좋지 않다고 알려져 있음
• 대부분의 현재 연구는 Width Pruning 또는 Width+Depth Pruning 방식에 초점을 맞추고 있으며, Inference Efficiency에 대한 Width Pruning과 Depth Pruning의 비교 분석은 거의 이루어지지 않음.

UnStructed / Semi-Structed Pruning

• Weight의 값을을 Sparse하게 0으로 만들어줌.
• 이렇게 zero-masking을 시키는 알고리즘은 성능이 좋지만, 특정 하드웨어 의존성이 매우 높음
• e.g) SparseGPT, Wanda

Structed Pruning

• 전체 layer나 MHA, FFN 등과 같은 패턴을 제거함
• 특정 하드웨어 의존성 문제 없이 Inference Efficiency 증대
• Structed Pruning은 Width Pruning (MHA Pruning), Depth Pruning (Layer Prunig), Width+Depth Pruning 이 존재함.
• e.g) Width Prunig - FLAP, LLM-Pruner
• e.g) Width+Depth Pruning - Sheared-LLaMA

Width Pruning 의 한계

• 제한된 환경에서 작은 batch size에 대해서는 Width Pruning이 효과가 없음

효과적인 Depth Pruning 방법 제안

• Block Pruning + LoRA Apaptor
• Depth Pruning이 Width Pruning과 비교하여 경쟁력이 있다는 것을 증명
• 작은 Batch size에서 Inference 속도 크게 개선

2. Method: Block Pruning

• Transformer Block 의 점수를 매겨 덜 중요한 블록들을 한번에 가지치기 후, LoRA Tuning 수행

Block 별 중요도 계산

• 블록 별 중요도 계산 방법에는 Mag, Taylor, Mag+, Taylor+, PPL 총 5가지 방법 존재.
  • MAG :
    단순히 가중치 절댓값들의 합이 크다면 중요한 블록, 값이 작다면 작은블록으로 판단
  • Taylor :
    가중치를 0으로 두었을 때, loss값이 크면 중요한 블록, loss값이 작으면 중요하지 않은 블록으로 판단
  • MAG+ & Taylor+ : 처음 4개 블록, 마지막 2개 블록을 Pruning 후보에서 제외하고 MAG, Taylor를 지표로 설정하였을 때
  • PPL (Perplexity) :
    모델이 주어진 데이터셋을 얼마나 잘 예측하는지를 평가하는 데 사용되는 일반적인 언어 모델링 지표
    수식 설명은 https://velog.io/@hyunku/NLP-Evaluation-Metric 참고

    Summary
    MAG : 가중치의 절댓값을 기준으로 블록의 중요도를 평가
    Taylor : 가중치가 모델 손실에 미치는 영향을 평가
    PPL : 모델이 주어진 데이터셋을 얼마나 잘 예측하는지를 평가하여 블록의 중요도를 판단

• 그 중 최종 Taylor+, PPL 선정
• 5가지 방법으로 블록 별 중요도를 계산하여 가지치기한 결과를 PPL, 상식 추론 정확도 2가지 평가지표로 평가함.
  
• 상식 추론 정확도는 Taylor+가 가장 좋았으며, PPL 성능은 PPL이 가장 좋았음
• 따라서 Taylor+와 PPL 두 가지 모두를 Block Pruning의 기준지표로 사용하였음.

일괄 가지치기

• 블록 수준 중요도 점수를 정렬한 후, 덜 중요한 블록을 한 번에 가지치기함

LoRA ReTraining

• Pruning된 모델을 효율적으로 재훈련하여 성능을 회복

3. Experimental Setup

Model

• LLaMA 7B
• Vicuna 7B, 13B

Baseline

• Width 기반 Pruning 기법들과의 비교를 위하여 Width 기반 알고리즘 채택
• FLAP, LLM-Pruner, Wanda-sp 3가지 Width Pruning 기법 선택

Dataset

• Pruning 단계에서 BookCorpus 으로부터 무작위로 10개의 샘플을 선택하여 블록 중요성을 계산.
• LoRA ReTraining에서는 Alpaca학습에 사용된 데이터의 50K 샘플을 사용함.

Evaluation

• lm-evaluation-harness 패키지의 7가지 상식 추론 데이터셋 사용 (Zero-shot)
• WikiText2, PTB에서의 PPL 성능 측정

Latency & Throughput (지연시간과 처리량)

• 배치 크기 M과 출력 시퀀스 길이 L(입력 길이 제외)이 주어질 때, Latency T는 prompt를 처리하여 ML 결과의 output token을 생성하는 데 필요한 시간을 의미함.
• Throughput은 ML/T로 계산. 초기 10개의 워밍업 배치 후 20번의 실행에서의 평균 결과를 기록함.

4. Result

Width Pruning vs Depth Pruning

• Width Pruning 방법은 Latency 개선으로 이어지지 않음.
• 이는 작은 배치에서는 계산 자체는 빨리 끝나지만, 메모리 access 빈도는 유지되므로 메모리 access 작업에서 병목 현상이 생겨 추론 속도를 제한함.
• Width Pruning 방법과 성능이 비슷함

Depth Pruning + Width Pruning

• 5B 보다 큰 모델에서는 MHA+FFN을 제거하는 Width Pruning이 더 좋았지만, 5B보다 작은 모델에서는 Layer를 제거하는 Depth Pruning이 더 좋았음.

One shot Pruning vs Iterative Pruning

• 한번에 가지치기하고 calibration 데이터셋으로 한번 학습시켰을 때의 성능과 하나의 블록을 제거하고 학습하고 다시 제거하고 학습하는 과정을 반복하며 가지치기를 하였을 때의 성능이 비슷함
• 한번에 가지치기하는 one shot pruning 방법이 훈련 비용면에서 매우 효율적

5. Conclusion

• Depth Pruning 방법이 Width Pruning 방법에 견줄만 하다.
• 특히 작은 배치 사이즈에서 큰 속도 개선이 있었다.