기술 요약: RAP (LLM 추론을 위한 런타임 적응형 가지치기)

문제 제기

대규모 언어 모델 (LLM) 은 막대한 계산 및 메모리 발자국으로 인해 배포에 상당한 장벽에 직면해 있습니다. 모델 압축 기술 (가지치기, 양자화, 증류) 은 해결책을 제시하지만, 기존 방법들은 다음과 같은 치명적인 한계를 겪고 있습니다:

1. 정적 휴리스틱: 현재 대부분의 가지치기 전략은 정적 워크로드를 가정하는 고정된 일회성 휴리스틱에 의존합니다. 이는 변동하는 사용자 요청 (예: 다양한 배치 크기와 시퀀스 길이) 으로 인한 런타임 메모리 변동을 적응하지 못합니다.
2. KV 캐시 무시: 전통적인 가지치기는 주로 정적 모델 가중치 (파라미터) 감축에 집중합니다. 그러나 자기회귀 추론에서 키 - 값 (KV) 캐시 메모리는 배치 크기와 시퀀스 길이에 따라 선형적으로 증가하여 종종 정적 파라미터 메모리를 압도합니다. 고정된 가지치기 비율은 파라미터 중심에서 KV 캐시 중심의 메모리 병목 현상으로의 동적 전환을 수용할 수 없습니다.
3. 시스템 수준 변동성: 실제 세계의 엣지 및 클라우드 환경은 함께 실행되는 애플리케이션과 간섭으로 인해 확률적 메모리 가용성을 경험하며, 이는 고정된 예산 가정을 무효화합니다.

해결된 핵심 과제는 출력 품질을 저하시키지 않으면서 변동하는 메모리 예산을 충족하면서 이질적이고 시간에 따라 변하는 요청 워크로드에 동적으로 적응하는 최적의 LLM 가지치기 정책을 선택하는 방법입니다.

방법론: RAP 프레임워크

저자들은 Reinforcement Learning (RL) 에 기반하여 실시간 시스템 상태에 따라 압축 전략을 동적으로 조정하는 탄력적 가지치기 프레임워크인 RAP (런타임 적응형 가지치기) 를 제안합니다.

1. 탐욕적 순차적 중요도 (GSI)

"일회성" 가지치기 (레이어 간 종속성을 무시하고 성능 저하를 초래함) 의 한계를 해결하기 위해, RAP 는 탐욕적 순차적 중요도 (GSI) 를 사용합니다.

• 메커니즘: 모든 블록을 한 번에 점수 매기는 대신, GSI 는 모델 성능 (퍼플렉시티로 측정) 의 최소 악화를 초래하는 블록을 반복적으로 제거합니다.
• 재교정: 각 제거 후, 축소된 아키텍처 내에서 나머지 블록의 중요도 점수가 재평가됩니다. 이 단계별 과정은 오차 누적을 제어하고 모델 구조가 변경됨에 따라 레이어의 동적 민감도를 고려합니다.

2. 강화 학습 공식화

RAP 는 가지치기 정책 선택을 강화 학습 에이전트가 메모리 효율성과 모델 충실도 사이의 균형을 배우는 순차적 의사결정 문제 (마르코프 의사결정 과정) 로 재구성합니다.

• 상태 ($S_t$): 에이전트는 각 추론 단계 $t$에서 복합 상태를 관찰합니다:
  • 요청 특성: 배치 크기 ($R_{bs}$) 와 시퀀스 길이 ($R_{sql}$).
  • 모델 구성: GSI 에서 유도된 모든 멀티헤드 어텐션 (MHA) 과 피드포워드 네트워크 (FFN) 블록의 중요도 점수.
  • 시스템 상태: 현재 사용 가능한 시스템 메모리 ($Sys_{avail}$) 와 제안된 행동에 대한 추정 메모리 오버헤드 ($Sys_{req}$).
• 행동 ($A_t$): 에이전트는 특정 MHA 및 FFN 블록을 유지할지 제거할지 결정하여 가지치기 정책을 선택합니다. 행동 공간 ($2^{2N}$) 의 조합적 폭발을 관리하기 위해, 결정은 순차적 단일 블록 선택으로 분해됩니다.
• 보상 ($R_t$): 통합 보상 함수는 두 가지 목표를 균형 있게 조정합니다:
  • 퍼플렉시티 ($R_{ppl}$): 언어 모델링 성능 유지를 장려합니다.
  • 메모리 페널티 ($R_{mem}$): 추론 중 피크 메모리 소비에 페널티를 부과합니다.
  • 보상은 가중 합입니다: $R_t = \sum (At)_i (\alpha R_{ppl, i} - \beta R_{mem, i})$, 여기서 하이퍼파라미터 $\alpha$와 $\beta$는 정확도 - 메모리 트레이드오프를 조정합니다.

3. 실행 흐름

1. RL 에이전트는 실시간 환경 (입력 길이, 배치, 메모리 가용성) 을 관찰합니다.
2. 현재 상태에 기반하여 가지치기 행동 (어떤 블록을 제거할지) 을 선택합니다.
3. 기본 모델이 가지치기된 아키텍처에서 추론을 실행합니다.
4. 에이전트는 결과적인 메모리 사용량과 퍼플렉시티에 기반한 보상을 받아 향후 의사결정을 최적화하기 위해 정책을 업데이트합니다.

주요 기여

1. 런타임 적응형 프레임워크: RAP 는 모델 가중치와 KV 캐시 메모리를 "실시간으로" 함께 고려하는 최초의 프레임워크로, 즉각적인 워크로드와 장치 상태에 기반하여 희소성 수준을 동적으로 조정합니다.
2. RL 기반 정책: 저자들은 가지치기를 RL 문제로 공식화하여, 메모리 절감과 생성 충실도 사이의 균형을 맞추는 최적 정책을 학습하도록 에이전트를 훈련시켰으며, 이는 정적 휴리스틱 베이스라인을 능가합니다.
3. 구조적 가지치기를 위한 GSI: 탐욕적 순차적 중요도 (GSI) 의 도입은 트랜스포머 가지치기에서의 레이어 간 종속성 문제를 해결하여, 공격적인 압축 하에서 블록 중요도를 결정하는 더 견고한 방법을 제공합니다.
4. 실증적 검증: 광범위한 실험은 RAP 가 다양한 모델 (Llama2, Llama3, Qwen) 과 메모리 제약 조건에서 최신 정적 가지치기 베이스라인 (예: LLM-Pruner, SliceGPT, ShortGPT) 일관적으로 능가함을 입증합니다.

실험 결과

실험은 80% 및 60% 통합 메모리 예산 (파라미터 및 KV 캐시 모두 포함) 하에서 Llama2-7B, Llama3-8B, Qwen1.5-7B, Qwen2.5-7B 에서 수행되었습니다.

• 생성 능력: RAP 는 가장 작은 퍼플렉시티 변동을 보여주었습니다. 80% 예산의 Llama2-7B 에서 RAP 는 WikiText2 에서 +6.33, PTB 에서 +21.47 만큼 퍼플렉시티를 증가시켰으며, 이는 각각 22.95 와 146.93 의 차이로 다음으로 좋은 방법보다 크게 우세했습니다.
• 하류 작업: 제로샷 상식 추론에서 RAP 는 80% 예산 하에서 밀집 모델 정확도의 86% 를 유지하여, Llama2-7B 의 경우 +8.2%, Llama3-8B 의 경우 +11.5% 만큼 주요 베이스라인을 능가했습니다.
• 복원력: 심각한 60% 메모리 제약 하에서 RAP 는 다른 방법들이 치명적인 저하를 겪은 반면, 밀집 모델 성능의 60% 이상을 유지한 유일한 방법이었습니다.
• 절대 연구:
  • GSI 제거 (정적 일회성 점수 사용) 는 평균 상식 정확도를 13.17% 감소시켰습니다.
  • RL 에이전트 제거 (무작위 블록 드롭 사용) 는 학습된 정책과 비교하여 퍼플렉시티가 급격히 증가 (+WikiText2 에서 +301.71) 하여 적응형 의사결정의 필요성을 강조했습니다.
• 오버헤드: RL 컨트롤러는 소형 아키텍처 (18K 파라미터) 로 인해 실행 시간 오버헤드를 무시할 수준 (<1% 지연 시간 증가) 으로 추가하며, 오프라인 GSI 계산은 배포 전체에 걸쳐 분산되는 일회성 비용입니다.

중요성 및 주장

이 논문은 RAP 가 정적 압축 기술과 동적 실제 배포 시나리오 간의 간극을 메운다고 주장합니다. "일률적" 휴리스틱에서 벗어나면서 RAP 는 다음과 같은 것을 가능하게 합니다:

• 확장 가능한 엣지 배포: 필요한 경우 상당한 가중치를 가지치기하여 장치에 맞출 수 있게 하면서, 자원이 허용될 때는 전체 용량을 유지함으로써, LLM 이 엄격하고 변동하는 메모리 제약 하에서 장치에 배포될 수 있게 합니다.
• 품질 손실 없는 효율성: 출력 품질에 미치는 영향은 최소화하면서 우수한 메모리 절감과 더 빠른 추론을 달성하여, 파라미터 중심과 KV 캐시 중심 병목 현상 사이의 트레이드오프를 효과적으로 탐색합니다.
• 일반화 능력: 이 프레임워크는 아키텍처에 구애받지 않는 것으로 입증되어, 다양한 모델 계열 (Llama, Qwen) 과 세대에서 잘 작동합니다.

저자들은 RAP 를 고정된 전략이 실패하는 이질적인 환경에서 LLM 추론을 실현 가능하게 만드는 유연하고 문맥 인식형 압축 메커니즘으로 위치시킵니다.