SpenseGPT: Practical One-shot Pruning Enabling Sparse and Dense GEMMs for LLM Inference

기술 요약: SpenseGPT: LLM 추론을 위한 희소 및 밀집 GEMM 대상 실용적 원샷 프루닝(One-shot Pruning)

1. 문제 정의

반구조적(Semi-structured) 2:4 희소성(4개의 연속된 가중치 중 2개가 0으로 프루닝됨)은 현대적인 가속기(예: NVIDIA B200, AMD GPU, Meta MTIA)에서 널리 지원되며, 하드웨어가 zero-weight 연산을 건너뛸 수 있게 함으로써 이론적으로 2배의 속도 향상을 제공합니다. 그러나 대규모 언어 모델(LLM) 추론에 있어 실질적인 도입에는 두 가지 주요 장애물이 존재합니다:

1. 엄격한 정확도 저하: 경직된 2:4 희소성 제약은 광범위한 재학습을 수행하지 않는 한, 특히 추론(reasoning) 및 지시 이행(instruction-following) 작업에서 심각한 정확도 손실을 초래할 수 있습니다.
2. 기존 완화 기법의 비실용성: 2:4 제약을 완화하려는 최근 연구들(예: PATCH, SlideSparse)은 결정적인 한계를 가집니다. PATCH는 아직 널리 보급되지 않은 특수 컴파일러(Triton) 지원이 필요하며, SlideSparse는 입력 활성화(activation) 확장을 필요로 하여, 특히 긴 시퀀스 길이($M$)에 대해 엔드 투 엔드(end-to-end) 속도 향상을 상쇄하는 런타임 오버헤드를 발생시킵니다.

2. 방법론

본 논문은 하이브리드 희소-밀집 포맷인 Spense와 이를 기반으로 구축된 원샷 사후 학습 프루닝(post-training pruning) 방법인 SpenseGPT를 제안합니다.

A. Spense 포맷
전체 가중치 행렬에 대해 2:4 희소성을 강제하는 대신, Spense는 각 가중치 행렬을 두 개의 연속된 영역으로 분할합니다:

• 밀집 영역(dense region): 가중치의 $p\%$를 유지합니다.
• 2:4 희소 영역(2:4 sparse region): 표준 2:4 희소성으로 프루닝됩니다.

이 설계는 커스텀 컴파일러나 입력 활성화 확장 없이도 기존의 고도로 최적화된 GEMM 라이브러리(예: NVIDIA GPU의 경우 sparse를 위한 cuSPARSELt 및 dense를 위한 cuBLASLt)를 활용합니다. 저자들은 두 가지 변형을 구현했습니다:

• 출력 분할(Output-split): 출력 텐서를 분할하여, 한 부분은 sparse GEMM을 사용하고 다른 부분은 dense GEMM을 사용합니다.
• 입력 분할(Input-split): 입력 텐서를 분할하여, 첫 번째 부분은 sparse GEMM으로 처리하고 두 번째 부분은 (누적 연산과 함께) dense GEMM으로 처리합니다.
  GLU 변형 MLP의 경우, 저자들은 $W_{up}$과 $W_{gate}$에는 output-split을 사용하고, $W_{down}$에는 input-split을 사용합니다.

B. 밀집 인덱스 선택 전략 (Dense Index Selection Strategies)
어떤 중간 인덱스를 밀집 상태로 유지할 것인가가 모델 품질에 결정적인 영향을 미친다는 점이 핵심 통찰입니다. 본 논문은 GLU 스타일 MLP의 중간 인덱스들이 수학적 출력을 변경하지 않고도 치환(permute)될 수 있음을 식별하였으며, 이를 통해 특정 인덱스를 밀집 상태로 선택한 후 실행을 위해 연속적인 블록으로 재배열할 수 있습니다. 두 가지 전략이 제안되었습니다:

1. SpenseGPT (단순 전략): 마지막 $p\%$의 중간 인덱스를 밀집 상태로 유지하도록 선택합니다. 이는 처리된 열(column)은 고정되고, 남은 오른쪽 열들로 재구성 오차를 보상하는 SparseGPT의 열 방향 처리 순서와 일치합니다. 밀집 열을 오른쪽에 둠으로써 "보상 버퍼(compensation buffer)" 역할을 하게 합니다.
2. SpenseGPT+ (손실 인식 전략): 추정된 재구성 손실(reconstruction loss)을 기반으로 밀집 인덱스를 선택합니다. Wanda에서 영감을 받아, gate, up, down 투영 전체에 걸친 기여도를 집계하여 각 중간 인덱스에 대한 중요도 점수를 계산합니다. 이 점수는 가중치 크기와 입력 활성화 크기를 결합합니다 ($S_{ij} = |W_{ij}|\sqrt{H_{j,j}}$). 점수가 높은 상위 인덱스들을 밀집 상태로 유지하도록 선택한 후, 남은 희소 영역에 SparseGPT를 적용하기 전에 이들을 연속적인 영역으로 치환합니다.

3. 주요 기여

• Spense 포맷: 기존의 고성능 희소 및 밀집 GEMM 라이브러리와의 호환성을 유지하면서 2:4 희소성 제약을 완화하는 하이브리드 희소-밀집 포맷을 제공하며, 커스텀 컴파일러 지원이나 입력 확장 오버헤드를 피합니다.
• 밀집 선택의 중요성: 밀집 인덱스를 선택하는 방식이 모델 품질 유지에 결정적인 요소임을 입증하였으며, 단순한 선택 방식은 심각한 성능 저하를 초래함을 보여주었습니다.
• 프루닝 전략: 원샷 프루닝을 위한 SpenseGPT(연속적 꼬리 선택) 및 SpenseGPT+(재구성 손실 기반 선택)를 도입하였습니다.
• 실제 속도 향상: 최근 GPU(B200)에서 반구조적 희소 텐서 코어를 사용하여 모델 품질을 유지하면서도, 원샷 프루닝이 실제 엔드 투 엔드 LLM 디코딩 속도 향상을 달성할 수 있음을 최초로 입증하였습니다.

4. 실험 결과

방법론들은 B200 GPU에서 FP8 정밀도를 사용하여 Qwen3-32B 및 Seed-OSS-36B 모델에 대해 평가되었습니다. 벤치마크에는 AIME (수학), GPQA (지식), IFEval (지시 이행), LiveCodeBench (코딩)가 포함되었습니다.

• 정확도 보존:
  • 엄격한 2:4 프루닝(50% 희소성)은 상당한 정확도 하락을 초래했습니다 (예: Seed-OSS-36B의 평균 점수가 SparseGPT 적용 시 73.86에서 65.18로 하락).
  • **SpenseGPT+**는 25% 희소성에서 Seed-OSS-36B에 대해 73.12의 평균 점수를 달성하여, 밀집 베이스라인(73.86)에 거의 근접했습니다.
  • Qwen3-32B의 경우, SpenseGPT+는 25% 희소성에서 70.76에 도달하여 밀집 베이스라인인 71.68에 근접했습니다.
  • 비교 실험 결과, SpenseGPT+는 밀집 인덱스를 올바르게 선택함으로써 SpenseGPT 및 다른 베이스라인(HyperPrune, Wanda)보다 성능이 월등히 뛰어남을 보여주었습니다.
• 엔드 투 엔드 속도 향상:
  • Spense는 B200 GPU에서 밀집 베이스라인 대비 최대 1.2배의 엔드 투 엔드 디코딩 속도 향상을 달성했습니다.
  • 반면, SlideSparse는 입력 활성화 확대로 인한 비-GEMM 오버헤드로 인해 유의미한 속도 향상을 제공하지 못했습니다.
• 효율성:
  • 32B–36B 모델을 프루닝하는 데 단일 B200 GPU에서 약 1.33–2.15시간이 소요되었으며, 이는 SparseGPT와 유사하며 HyperPrune보다 훨씬 빠릅니다. ProxSparse는 이 모델들에서 메모리 부족(OOM) 문제를 겪었습니다.

5. 의의 및 주장

본 논문은 최근 하드웨어(특히 B200 GPU)의 반구조적 희소 텐서 코어로부터 파생된, 모델 품질을 유지하면서도 실제 엔드 투 엔드 LLM 디코딩 속도 향상을 달성한 최초의 원샷 프루닝 시연이라고 주장합니다.

저자들은 자신들의 접근 방식이 이론적인 하드웨어 가속과 실제 배포 사이의 간극을 메운다는 점을 강조합니다. 특수 컴파일러나 입력 확장이 필요하지 않으므로, Spense는 기존의 최적화된 라이브러리를 즉시 활용할 수 있게 합니다. 이 연구는 단순히 희소성 제약을 완화하는 것만으로는 부족하며, 어떤 가중치를 밀집 상태로 유지할지에 대한 전략이 모델의 추론 및 지시 이행 능력을 보존하는 데 매우 중요하다는 점을 부각합니다.

인정된 한계점:

• 속도 향상은 기초가 되는 GEMM 라이브러리의 효율성에 따라 달라지며, 배치 크기, 시퀀스 길이 및 하드웨어에 따라 차이가 있습니다.
• SpenseGPT+는 캘리브레이션 데이터에 의존합니다. 만약 배포 시의 분포가 캘리브레이션 세트와 크게 다를 경우, 선택된 밀집 인덱스가 최적이 아닐 수 있습니다.