Reducing the GPU Memory Bottleneck with Lossless Compression for ML -- Extended

기술 요약: ML을 위한 무손실 압축을 통한 GPU 메모리 병목 현상 완화

문제 정의

그래프 신경망(GNN), 딥러닝 추천 모델(DLRM), 대규모 언어 모델(LLM)을 포함한 현대의 머신러닝(ML) 워크로드는 수백 기가바이트에서 테라바이트에 이르는 데이터셋을 빈번하게 처리합니다. 이러한 데이터 볼륨은 가용 GPU 메모리 용량을 초과하며, 이에 따라 PCIe 인터커넥트를 통해 CPU 메모리, 디스크 또는 네트워크 스토리지로부터 온디맨드 데이터 전송이 필요하게 됩니다.

이러한 시나리오에서의 주요 성능 제약 요인은 PCIe 대역폭이며, 이는 GPU 메모리 대역폭보다 한 자릿수 이상 낮습니다. 결과적으로 실행은 GPU 연산 능력보다는 데이터 전송 레이턴시에 의해 병목 현상이 발생합니다. 캐싱 및 연산과 전송의 오버랩(overlapping) 기술이 부분적인 완화책을 제공하지만, 특히 캐시 미스가 발생할 때 이를 완전히 해결할 수는 없습니다.

기존 솔루션은 데이터 볼륨을 줄이기 위해 양자화(quantization)와 같은 **손실 압축(lossy compression)**에 의존하는 경우가 많습니다. 그러나 손실 방식은 모델과 워크로드에 따라 변동성이 큰 예측 불가능한 정확도 저하를 초래합니다. 상용 배포 환경에서는 미세한 정확도 손실(예: 0.1%)조차 수익 리스크로 인해 수용 불가능한 경우가 많으므로, 손실 압축을 범용적인 솔루션으로 채택하기는 어렵습니다.

**무손실 압축(Lossless compression)**은 모델 정확도를 보존하면서도 배포 친화적인 대안을 제공하지만, ML 파이프라인 내에서 다음과 같은 중대한 구현 과제에 직면합니다:

1. 레이턴시 경쟁: CPU-GPU 전송 레이턴시는 GPU 압축 해제 레이턴시와 경쟁 관계에 있습니다. 전통적인 압축 알고리즘은 종종 전송 절감 효과를 상쇄하는 오버헤드를 추가합니다.
2. 메모리 제약: 압축 해제 과정은 흔치 않게 대규모 중간 버퍼나 복잡한 메타데이터 구조를 요구하며, 이는 희소한 GPU 메모리를 두고 모델 입력값과 경쟁하게 됩니다.
3. 병렬성 불일치: 많은 표준 무손실 알고리즘(예: Huffman coding, LZ4)은 딕셔너리 룩업이나 불규칙한 메모리 액세스 패턴에 의존하여 GPU 워프(warp) 병렬성과 병합된(coalesced) 메모리 액세스를 저해합니다.

방법론: 불변 비트 패킹 (Invariant Bit Packing, IBP)

저자들은 ML 텐서 전송을 위해 특별히 설계된 새로운 무손실 압축 알고리즘인 **불변 비트 패킹(IBP)**을 제안합니다. 핵심 통찰은 ML 텐서들이 수치적으로 다양함에도 불구하고, 텐서 그룹 간에 **저엔트로피 불변 비트(low-entropy invariant bits)**를 공유하는 경우가 많다는 점입니다. 예를 들어, 부동 소수점 표현에서 부호(sign) 및 지수(exponent) 비트는 정규화 및 학습 역학으로 인해 많은 텐서에 걸쳐 일정하게 유지되는 경우가 많습니다.

핵심 알고리즘

IBP는 세 가지 주요 단계로 작동합니다:

1. 전처리 (메타데이터 생성):

   • 시스템은 입력 텐서 배치를 분석하여 높은 비율($T$)의 텐서들 사이에서 불변(즉, 값이 동일한)인 비트 위치를 식별합니다.
   • 시스템은 마스크(Mask)(어떤 비트 위치가 불변인지 나타냄)와 비트값(Bitval)(불변 값을 저장)을 생성합니다.
   • 이 메타데이터는 GPU 메모리에 저장되며(통상 킬로바이트 단위), 전체 데이터셋에 재사용되어 텐서당 메타데이터 오버헤드를 방지합니다.

2. 압축:

   • 입력 텐서는 청크(chunk) 단위로 처리됩니다. 각 청크에 대해 알고리즘은 불변 위치의 비트가 저장된 Bitval과 일치하는지 확인합니다.
   • 일치하는 경우, 불변 비트가 제거(elimination)되고 나머지 비트들이 빈 공간을 채우기 위해 이동(shift)됩니다.
   • 각 청크에는 압축되었는지 아니면 압축되지 않은 상태로 남겨졌는지(압축이 공간 절감 효과를 내지 못하는 경우를 처리하기 위함)를 나타내는 **참여 비트(participation bit)**가 앞에 붙습니다.
   • 압축은 제로 카피(zero-copy) 통합 가상 메모리(Unified Virtual Memory)를 사용하여 GPU에서 수행되며, 이를 통해 GPU가 명시적인 CPU 주도 DMA 설정 없이 CPU 메모리에 직접 접근할 수 있게 합니다.

3. 압축 해제:

   • 압축 해제는 GPU 워프 병렬성에 최적화되었습니다. 32개의 스레드로 구성된 하나의 워프가 협력하여 CPU 메모리로부터 연속적인 압축 데이터 세그먼트를 공유 메모리 작업 공간으로 읽어옵니다.
   • 스레드들은 워프 프리미티브(warp primitives)(워프 내 저지연 통신)를 사용하여 접두사 스캔(prefix scan)을 수행하고, 각 스레드에 할당된 청크의 시작 비트 오프셋을 계산합니다.
   • 그 후, 스레드들은 글로벌 Bitval로부터 불변 비트를 올바른 위치에 다시 삽입하여 데이터를 독립적으로 재구성합니다.
   • 이 프로세스는 데이터 페칭(fetching)과 압축 해제 로직을 중첩시켜 PCIe 레이턴시를 숨기기 위해 비동기 메모리 액세스를 활용합니다.

시스템 통합

저자들은 IBP를 두 가지 주요 방식으로 ML 파이프라인에 통합합니다:

• 캐시 압축: GPU 소프트웨어 캐시(예: GNN 프레임워크 내)에 저장된 텐서를 압축하여 유효 캐시 용량을 증대시킵니다.
• PCIE 전송 압축: 텐서를 CPU 메모리에서 전송하기 전에 압축합니다. 가변 크기의 압축 텐서를 처리하기 위해, 시스템은 인플레이스 압축(in-place compression)(고정 크기 버퍼를 할당하되 압축된 부분만 읽음)을 사용하며, 어떤 텐서가 압축되었는지 추적하기 위해 GPU 메모리에 **비트마스크(bitmask)**를 유지합니다.

주요 기여

1. ML 데이터 불변성 분석: 본 논문은 다양한 ML 워크로드(GNN, DLRM, LLM) 및 데이터 타입(FP32, FP16, BF16)에서 불변 비트의 존재와 양을 정량화하여, 구조화된 분포가 정확도 손실 없이 상당한 무손실 압축을 가능하게 함을 입증했습니다.
2. 불변 비트 패킹(IBP) 알고리즘: 불변 비트를 제거하고 최소한의 글로벌 메타데이터에 저장하며, 워프 병렬(warp-parallel) 압축 해제를 활용하여 높은 처리량을 달ing하는 GPU 최적화 무손실 압축 기술을 제시합니다.
3. 시스템 통합: IBP를 사용하기 쉬운 API(PyTorch 확장 및 CUDA 라이브러리)를 통해 인기 있는 ML 프레임워크(DGL, Legion, Colossal-AI, FlexGen, InfiniGen)에 통합했습니다.
4. 포괄적 평가: 최첨단 GPU 가속 라이브러리(nvCOMP, ndzip-gpu) 및 네이티브 시스템과 실제 데이터셋을 대상으로 광범위한 벤치마킹을 수행했습니다.

결과

평가는 NVIDIA A100 GPU와 PCIe Gen 4 환경에서 수행되었습니다.

• 압축 성능: IBP는 조밀한(dense) GNN 데이터셋에서 10–12%, DLRM 가중치에서 8%, 그리고 LLM 가중치 및 KV-캐시(특히 BF16 데이터 타입)에서 최대 **27%**의 공간 절감 효과를 달성했습니다.
• 처리량: IBP는 기존의 GPU 가속 압축 라이브러리들을 지속적으로 능가합니다. 희소(sparse) GNN 데이터셋의 경우, 압축되지 않은 전송 대비 9.7배 빠른 전송 처리량을 달성했습니다.
• 워크로드 가속:
  • GNN 학습: PCIe 전송 시간을 줄이고 캐시 적중률을 높임으로써 평균 **74%**의 가속(특정 사례에서는 최대 2.7배)을 달성했습니다.
  • DLRM 추론: 효율적인 GPU 주도 전송 및 CPU 오버헤드 감소를 통해 임베딩 룩업에서 평균 **180%**의 가속을 달성했습니다.
  • LLM 추론: 오프로딩된 추론에서 **24%**의 가속을 달성했습니다. 가중치 오프로딩(BF16)의 경우 레이턴시가 22% 감소했으며, KV-캐시 오프로딩의 경우 레이턴시가 최대 27%까지 감소했습니다.
• 오버헤드: 압축 해제 오버헤드는 미미합니다. 공간 절감률이 50% 미만인 경우(일반적인 조밀한 ML 데이터의 경우), 처리량은 이상적인 압축되지 않은 전송 속도의 95% 이상을 유지합니다. 전처리 시간은 보통 수준이며 크리티컬 패스(critical path) 외부에서 발생합니다.

의의 및 주장

본 논문은 IBP가 양자화나 손실 압축과 관련된 정확도 리스크 없이 GPU 메모리 병목 현상을 해결할 수 있는 실용적인 무손실 솔루션을 제공한다고 주장합니다. ML 데이터의 고유한 구조(불변 비트)를 활용하고 GPU 아키텍처(워프 병렬성, 병합된 액세스, 제로 카피)에 최적화함으로써, IBP는 압축 해제 오버헤드를 PCIe 전송 레이턴시 뒤로 효과적으로 숨깁니다.

저자들은 IBP가 다양한 데이터 타입(float32, float16, bfloat16, 희소 벡터) 및 워크로드(학습 및 추론)에 적용 가능한 범용적 알고리즘임을 강조합니다. 기존 프레임워크에 단 35줄 미만의 파이썬 코드로 IBP를 통합할 수 있다는 점은, 성능과 정확도를 유지하면서 GPU 메모리 한계를 넘어 모델을 확장하려는 상용 ML 배포 환경에서의 채택 가능성을 보여줍니다.