Shortened LLaMA: Depth Pruning for Large Language Models with Comparison of Retraining Methods

ICLR 2024 Workshop

👀 요약 👀


✨ Point ✨
block-level (Transformer block)로 importance 계산 (mag / taylor / ppl)
LoRA / Continued Pretraining으로 재학습 (retraining)

 

Abstract로 흐름 파악하기

pruning의 방법으로 width와 depth가 있는데,

대부분의 연구가 width-only 또는 blend of width and depth로 이루어지고 있다. 

 

본 연구에서는 간단한 depth pruning만으로도 LLM을 효과적으로 compress할 수 있다고(?) 주장한다.

추론 속도도 빠르며, 메모리 제약도 고려하였다(제한된 Batch size)

---

 

1. Introduction

- 본 연구는 structured pruning에 관한다. 

structured pruning은 불필요한 가중치 그룹을 제거하고, hardware-agnostic acceleration을 용이하게 한다.

* hardware-agnostic acceleration: 특정 하드웨어에 종속되지 않고 다양한 하드웨어(예: CPU, GPU, NPU, FPGA 등) 위에서도 잘 동작하면서 성능을 높이는 기술

 

- LLM inference는 autoregressive decoding 매커니즘을 따른다. 앞 토큰을 기반으로 토큰 하나씩 예측하는 방법이다.

이러한 방법은 memory-bound nature를 보이는데, 이는 GPU 연산능력 측면에서 상당한 비효율적이다.

-> batch size를 키우는 방법을 쓰고 있지만, 제한된 GPU 환경에서는 배치크기를 작게할 수밖에 없기 때문에, 이런 환경에서도 추론 속도를 높이고자 하였음.

 

- depth pruning은 큰 유닛을 제거하는 방법이다보니 width에 비해 비효율적이라고 여겨졌지만, 본 논문에서 꼭 그렇지는 않다는 것을 밝혔다. 

 

Contribution:

1. 제한된 배치사이즈에서, width pruning은 추론속도 향상 X

2. 간단하지만 효과적인 depth pruning 방법 제안

3. pruning ratio가 적당하면 LoRA를 통한 retraining이, ratio가 커지면 full-parameter update가 성능 향상에 중요하다. 

 

2. Problem: Small-batch LLM Inference

 

our focus is on accelerating the inference of LLMs under small-batch conditions caused by hardware restrictions. (배치사이즈를 키워서 inference 속도를 향상시키는 게 아니라, pruning을 통해서 속도를 향상시킨다는 의미에서 언급한 것 같음. 배치사이즈를 키우는 방법은 GPU 부자일 때만 가능하므로, 작은 배치사이즈에서도 적용할 수 잇는 방법이다.-)

1. width 프루닝은 생성 속도를 향상시키지 않으며, 가중치의 차원이 GPU에 적합하지 않게 변하는 경우 성능이 저하되기도 한다.

2. 의미있는 속도 개선은 depth 프루닝을 통해서만 가능하다.

 

3. Method: Block Pruning

트랜스포머 블럭 자체를 프루닝할 하나의 유닛으로 본다.

방법: 간단한 metric으로 중요하지 않은 블럭을 식별하고, one-shot pruning을 한다.

3.1. Evaluation of Block-level Importance

: linear weight matrix 

size: (d_out, d_in)
k: type of operation (e.g. 멀티헤드어텐션의 query projection,, FFN의 up projection 등)
n: n번째 트랜스포머 블럭

 

- output neuraon level로 weight importance score을 계산했다...(뭐라는거)

 

Magnitude(Mag). 

작은 norm을 가진 weight은 덜 중요한 정보이다.

Taylor

기본 식

: 어떤 weight $W_{i,j}^{k,n}$를 0으로 만들었을 때 모델 성능(Loss)이 얼마나 변하는지(즉, 제거했을 때의 영향력)는,

그 weight의 크기와 손실에 대한 gradient의 곱의 절댓값으로 근사할 수 있다.

우항 : 변화를 1차 도함수로 근사한 값 (Taylor expansion)

특정 파라미터를 제거했을 때 생기는 error를 통해 중요도를 측정한다.

 

 

Mag+ and Taylor+

트랜스포머 초반 블럭은 중요하지 않다고 레이블되긴 하는데, 막상 없애면 성능 저하된다는 이전 연구 결과

-> 첫 4블럭, 마지막 2블럭은 보존

 

Perplexity (PPL)

각 블럭을 없애면서 PPL 변화를 측정 (calibration set 사용)

세타^n : n번째 블럭을 뺀 모델
s = 1,...,S : 시퀀스
l = 1,...,L : 토큰

- PPL은 next-token prediction loss에서 유도되고, forward-apss computation만 요구한다. 

앞 뒤 블럭을 빼면 ppl이 치솟는다..

(오................... 오... 오히려 중간 레이어를 뺀 경우가 ppl 변화가 적다.. 오..

 앞 뒤 블럭에 중요한 정보가 담겨잇는 건가??? 잉? 

그냥 단순히 '제거'만 한 거라서 ppl이 증가하는 거겟지? distill처럼 조치를 취하면 ㄱㅊ아질 거 같기두 ..하고.....)

accuracy 결과도 비슷하려나???

 

 

3.2. One-shot Pruning

블럭별로 중요도 순서를 구해두고 이제 프루닝을 진행할 순서이다.

블럭의 파라미터 개수를 구할 수 있으므로, 원하는 사이즈로 프루닝할 수 있다.

 

iterative pruning은 one-shot pruning보다 컴퓨팅 타임이 길다는 단점이 있다.

게다가 어떤 프루닝 scheme을 썼느냐보다 retraining 전략이 더 중요하다는 것을 관찰했다.

 

3.3. Retraining for Performance Restoration

structured pruning은 재학습이 필요없거나 낮은 재학습 비용으로 실현 가능하다고 최근 연구에서는 암시한다.

하지만 재학습 '방법 types'에 대한 분석은 충분하지 않기 때문에 진행하였다.

근데 프루닝을 하고 재학습을 하면.............................................좋은 초기화말고 장점이 더 있나? 흐으으음 

큰 모델을 어케 살릴 건지?? 생각하다가 MoE로 넘어갓다 옴 ,, 큰 모델과 전문가.........................

Low-Rank Adaptation (LoRA)

width에 적용한 이전 연구(Ma et al. (2023)) (리뷰) 를 따라 본 연구에도 적용해봄.

효과가 있었으나 프루닝 비율이 높아지면 제대로 작동하지 않는 현상 (width, depth 모두)

Continued Pretraining (CPT)

large-scale pretraining corpus로 모든 파라미터를 업데이트하는 방법.

LoRA보다는 더 많은 리소스가 필요하지만 학습 속도도 빠르고, 랜덤 초기화보다 더 우수한 성능을 낸다.

CPT -> LoRA

CPT 하고 나서 instruction set을 사용해서 LoRA를 적용해서 추가적인 성능향상이 있는지 보았다. 

 

4. Experimental Setup

Source Model

LLaMA-7B

Vicuna-{7B, 13B}-v1.3

Baseline

[ Width pruning ]

LLM-Pruner

FLAP

Wanda-sp

[ Retraining-free block pruning method ]

SLEB

 

 

Data

BookCorpus

Alpaca (for LoRA)

SlimPajama (for CPT)

Evaluation

zero-shot accuracy (BoolQ, PIQA, HellaSwag, WinoGrande, ARC-easy, ARC-challenge, OpenbookQA)

zero-shot PPL (WidiText2, PTB)

Latency and Throughput(처리량)

(추론 속도가 빠르다는 걸 보여주기 위함)

batch size: M

output sequence length: L

latency: T (M L 아웃풋 토큰 생성까지)

throughput: M L / T

 

Implementation

..

5. Results

5.1. Moderate Pruning and LoRA Retraining

- width pruning이 LLM 추론 효율성을 증가시키지 않는다.

- width pruning의 경우 속도가 오히려 증가한 경우도 있는데, GPU에 최적화되지 않은 차원으로 바뀌었기 때문이다 (e.g., FFN의 hidden size가 8로 안 나눠짐)

- 반면 depth는 속도도 빨라지고, 적은 메모리 접근을 가능하게 했다.

- LLM-Pruner도 똑같이 LoRA를 사용했지만 width랑 비슷하게 성능이 나왓다.. .. . 

- SLEB랑 비교했을 때, 사이즈가 작아지면 상대적으로 별로임

 

 

 

5.2. Aggressive Pruning and CPT Retraining

- 프루닝 ratio가 클 때(fewer than 3.7B param) LoRA-based tuning이랑 retraining-free approach 둘 다 성능이 낮다.

- CPT 방법은 효과적임 / CPT->LoRA는 zs accuracy는 조금 향상, PPL은 조금 악화

- CPT 방법은 LoRA랑 비교했을 때는 계산 비용 높긴한데, 기본 모델에 CPT(GPU8개로 2주)하는 것보다는 완전 효율적임(GPU1개로 하루)

 

- 60% 프루닝된 모델 (2.7B) , our는 잘 생성한다.

- 같은 사이즈의 모델인 경우, 랜덤 초기화보다 pruning으로 시작하는 게 더 좋은 결과

5.3. Applicability with Quantization

 

 

GPTQ 방법으로 큰 성능저하 없이 VRAM 사용을 줄였다. 

 

* PTQ: 모델 학습(traning)이 끝난 후, 무겁고 느린 모델을 가볍게(빠르게) 만들려고 숫자(파라미터, 연산 결과)를 ‘작게’ 바꿔주는 기술

* VRAM: Video RAM, GPU전용 메모리

5.4. Ablation Study

 

5.4.1. Importance Criteria for Block Pruning

'+' 표시가 없는 메서드들은 essential한 initial block들을 유지하는데 실패했음 -> 성능 저하

- 가중치의 크기에만 의존하는 Mag 방법보다는 Taylor 방법이 우수하다

 

5.4.2. Structural Unit for Depth Pruning

각 모듈(MHA, FFN)을 기준으로 제거했을 때의 영향을 측정 (+LoRA)

- 5B 보다 클 때는 각 모듈을 기준으로 제거했을 때 accuracy가 좀 더 높지만, 그 외는 블럭 단위가 더 나은 성능을 보임.

이는 작은 단위로 제거할수록 성능이 향상된다는 일반적인 믿음과 다르다. (엥? 당연히 .. 중간에 있는 MHA나 FFN이 제거되면 학습된 흐름이 끊기니까 엄청 안 좋을 거 같은디.;; 근데 그런거 치고는 성능이 높은 편인 것 같기도 하고)

 

- 사실 각 모듈의 공동의 역할이 있기 때문에 독립적으로 처리하는 게 최적의 방법은 아닐 수 있다고 언급한다.

     - Table 6에서 5.3B의 경우 일부 구간에서 FFN만 연속적으로 남은 경우도 있었다 -> attention이 없어져서 word interaction을 다루는 모델 능력 떨어졌을 것

     - 반면 블럭단위 프루닝은, 이웃한 블럭들이 손실된 정보와 유사한 기능을 했을 것

 

5.4.3. Calibration(교정) Data Volume

- block-level importance를 구하기 위해 calibration data를 사용했다. 

- Table 7을 보면 알 수 있듯 10개로도 충분하다.

- Taylor+의 경우에는, 1k개의 샘플을 쓸 때 정확도가 더 낮아지는 모습을 보인다.

(원인에 대한 탐구는 future research로 넘긴다)

6. Related Work

* SparseGPT (Frantar and Alistarh, 2023) addresses the layer-wise reconstruction problem for pruning by computing Hessian inverses.

* structured pruning removes organized patterns, such as layers (Fan et al., 2020; Jha et al., 2023), (아니 abstract만 보긴 햇지만 레이어 단위 아닌거 같은데 ;;;;;;) FFN’s hidden sizes (Nova et al., 2023; Santacroce et al., 2023), and some hybrid forms (Lagunas et al., 2021; Xia et al., 2022; Kwon et al., 2022; Kurtic et al., 2023)

* Sheared-LLaMA (Xia et al., 2024) introduces a mask learning phase aimed at identifying prunable components in both the network’s width and depth.

* depth pruning approaches (Song et al., 2024(SELB); Men et al., 2024(ShortGPT); Tang et al., 2024(Rethinking...))