[llm-compressor] Disable LM Head - 불필요한 LM Head 연산 비활성화

PR 링크: vllm-project/llm-compressor#2081 상태: Merged | 변경: +74/-10

들어가며

LLM의 LM Head(언어 모델 헤드)는 hidden state를 vocabulary 크기의 logit으로 변환하는 마지막 레이어다. vocabulary가 크면(100K+) 이 연산이 상당한 시간과 메모리를 소비한다. 양자화 캘리브레이션에서는 LM Head의 출력이 필요 없는 경우가 많은데, 이 PR은 이를 비활성화하는 옵션을 추가한다.

핵심 코드 분석

LM Head 비활성화의 효과:

기존: hidden_states -> [LM Head: hidden_dim x vocab_size GEMM] -> logits (사용 안함)
개선: hidden_states -> (LM Head 스킵) -> 캘리브레이션 완료

대부분의 양자화 기법(GPTQ, AWQ, SmoothQuant 등)은 중간 레이어의 활성화 분포만 필요하고 최종 logit은 사용하지 않는다. LM Head를 스킵하면:

• vocab_size 100K, hidden_dim 4096인 경우: 약 400M 파라미터의 GEMM 연산 절약
• 배치마다 절약: 캘리브레이션 데이터 전체에 대해 누적 효과

왜 이게 좋은가

1. 캘리브레이션 속도 향상: vocab이 큰 모델(Qwen: 151K)에서 특히 효과적
2. GPU 메모리 절약: LM Head의 활성화 메모리가 해제되어 더 큰 배치 가능
3. 정확도 영향 없음: 양자화 결과에 전혀 영향을 주지 않는 순수 최적화
4. 선택적 적용: 필요한 경우(perplexity 기반 양자화 등)에는 LM Head 유지 가능

정리

74줄의 변경으로 양자화 파이프라인에서 불필요한 연산을 제거하는 깔끔한 최적화다. 특히 대규모 vocabulary를 가진 최신 모델에서 체감되는 시간 절약을 제공한다.

참고 자료

• PR #2081

---

이 글은 AI(Claude)의 도움을 받아 작성되었습니다. 코드 분석 내용은 실제 PR diff를 기반으로 합니다.