본문으로 건너뛰기

[onnxruntime] ONNX Runtime CUDA EP, 2-bit 양자화 GEMM/GEMV 지원 추가로 모델 경량화 가속

PR 링크: microsoft/onnxruntime#32699 상태: Merged | 변경: +1321 / -153

들어가며

최근 딥러닝 모델의 크기가 기하급수적으로 증가하면서, 모델 경량화 및 추론 성능 최적화는 더 이상 선택이 아닌 필수가 되었습니다. 특히 대규모 언어 모델(LLM)과 같이 방대한 파라미터를 가진 모델에서는 메모리 대역폭과 연산량이 병목 현상의 주된 원인이 됩니다.

Microsoft의 ONNX Runtime은 다양한 하드웨어 가속기를 지원하는 강력한 추론 엔진으로, 지속적인 최적화 작업을 통해 모델의 효율성을 높이고 있습니다. 이번 PR(#XXXX)은 ONNX Runtime의 CUDA Provider에 2-bit 양자화된 가중치에 대한 GEMM(General Matrix Multiply) 및 GEMV(General Matrix-Vector Multiply) 연산 지원을 추가하여, 특히 LLM과 같이 MatMul 연산 비중이 높은 모델의 성능을 획기적으로 개선하는 것을 목표로 합니다.

이 글에서는 해당 PR의 코드 변경 사항을 분석하고, 왜 이러한 변경이 성능 향상에 기여하는지, 그리고 어떤 기술적 교훈을 얻을 수 있는지 심도 있게 살펴보겠습니다.

코드 분석

이번 PR은 주로 CUDA Provider의 양자화 관련 커널 및 빌드 설정을 수정하여 2-bit 가중치 지원을 추가하는 데 초점을 맞추고 있습니다.

1. CMakeLists.txt: 빌드 옵션 확장

빌드 시스템 설정 파일인 CMakeLists.txt에서는 새로운 기능 활성화를 위한 옵션들이 추가되거나 수정되었습니다.

Before:

cmake_dependent_option(onnxruntime_USE_FPA_INTB_GEMM_FULL
  "Build all FpA IntB GEMM CUDA kernel variants instead of the compact FP16/BF16 INT4/INT8 set" OFF
  "onnxruntime_USE_CUDA;onnxruntime_USE_FPA_INTB_GEMM" OFF)

After:

cmake_dependent_option(onnxruntime_USE_FPA_INTB_GEMM_FULL
  "Build all FpA IntB GEMM CUDA kernel variants instead of the compact FP16/BF16 INT2/INT4/INT8 set" OFF
  "onnxruntime_USE_CUDA;onnxruntime_USE_FPA_INTB_GEMM" OFF)

onnxruntime_USE_FPA_INTB_GEMM_FULL 옵션의 설명이 INT4/INT8에서 INT2/INT4/INT8로 확장되었습니다. 이는 기존의 4-bit 및 8-bit 지원에 더해 2-bit 양자화 커널까지 포함하도록 빌드 옵션의 범위를 넓혔음을 의미합니다. 또한, 빌드 시 출력되는 메시지 역시 INT4/INT8에서 INT2/INT4/INT8로 업데이트되어, 사용자가 어떤 커널 세트가 빌드되는지 명확히 인지할 수 있도록 했습니다.

2. onnxruntime_cuda_source_filters.cmake: LLM 커널 소스 필터링

LLM 관련 CUDA 커널 소스 파일을 필터링하는 로직에도 2-bit 지원이 반영되었습니다.

Before:

if(onnxruntime_USE_FPA_INTB_GEMM AND NOT onnxruntime_USE_FPA_INTB_GEMM_FULL AND
 ((_src MATCHES "/fpA_intB_gemm/" AND
 NOT _src MATCHES "/fpA_intB_gemm/(fp16|bf16)_int(4|8)_gemm_scaleonly\.cu$") OR
 (_src MATCHES "/fpA_intB_gemv/dispatcher_" AND
 NOT _src MATCHES "/fpA_intB_gemv/dispatcher_(fp16|bf16)_int(4|8)\.cu$")))

After:

if(onnxruntime_USE_FPA_INTB_GEMM AND NOT onnxruntime_USE_FPA_INTB_GEMM_FULL AND
 ((_src MATCHES "/fpA_intB_gemm/" AND
 NOT _src MATCHES "/fpA_intB_gemm/(fp16|bf16)_int(2|4|8)_gemm_scaleonly\.cu$") OR
 (_src MATCHES "/fpA_intB_gemv/dispatcher_" AND
 NOT _src MATCHES "/fpA_intB_gemv/dispatcher_(fp16|bf16)_int(2|4|8)\.cu$")))

onnxruntime_USE_FPA_INTB_GEMM_FULL 옵션이 비활성화된 경우, 즉 'compact' 빌드 모드일 때 포함될 커널을 정의하는 부분입니다. 기존에는 int(4|8)만 제외되었으나, 수정 후에는 int(2|4|8)로 변경되어 compact 빌드에서도 2-bit scale-only GEMM/GEMV 커널이 포함될 수 있음을 나타냅니다. 이는 기본 빌드 설정에서 2-bit 연산을 더 쉽게 활용할 수 있게 해줍니다.

3. matmul_nbits.md: 문서 업데이트

문서 파일 matmul_nbits.md에서도 2-bit 지원에 대한 설명이 추가 및 수정되었습니다.

Before:

- `bits == 2` → `TryMatMul2Bits` (no fused bias support; returns `false` if bias set).

After:

- `bits == 2` → `TryMatMul2Bits` (`1 <= M <= 8`; no fused bias support; larger `M` or bias falls through to §5).

2-bit 연산의 경우, 기존에는 bias 지원 여부만 명시되었으나, 수정 후에는 1 <= M <= 8이라는 제한 조건이 추가되었습니다. 이는 2-bit 커널이 특히 작은 배치 크기(M)에 최적화되어 있음을 시사하며, 더 큰 M이나 bias가 필요한 경우에는 기존의 dequantize + cuBLAS fallback 경로로 처리됨을 명확히 합니다.

또한, 문서의 여러 섹션에서 compact 빌드와 full 빌드의 차이점, 2-bit 가중치에 대한 제약 조건 (예: N의 배수 조건, block_size 요구사항), SM90 (Hopper) 아키텍처에서의 제약 사항 등이 상세하게 기술되었습니다.

특히, 2-bit 가중치에 대한 N의 정렬 요구사항이 128로 증가했으며 (bits==4의 64보다 높음), block_size는 64 또는 128만 지원된다는 점이 명시되었습니다. 이는 2-bit 데이터의 특성상 더 큰 메모리 접근 단위가 필요하기 때문입니다.

4. 내부 커널 로직 수정 (matmul_2bits.cu, dequantize_blockwise_2bits.cu 등)

PR 설명에 따르면, 2-bit 연산을 위한 새로운 커널 구현 및 기존 커널 수정이 이루어졌습니다. 주요 변경 사항은 다음과 같습니다:

  • INT2 수치 변환, SM80 인터리브 레이아웃 지원: 2-bit 정수 가중치를 FP16/BF16 활성화와 함께 처리하기 위한 변환 로직 및 SM80 아키텍처에 최적화된 가중치 레이아웃 지원이 추가되었습니다.
  • B-fragment 파이프라인 수정: 2-bit 연산에서 발생하는 odd single-load 케이스에 대한 B-fragment 파이프라인 문제를 해결하기 위해 레지스터 버퍼를 회전시키는 기법이 도입되었습니다. 이는 런타임 커서로 인한 추가적인 로컬 메모리 사용을 방지하고 성능 저하를 막습니다. (PR 설명 참조: "Fixes the two-deep B-fragment pipeline for the odd single-load INT2 case by rotating register buffers with compile-time indices.")
  • 정확도 테스트 강화: A = I 레이아웃 라운드 트립 테스트가 2-bit 및 4-bit 가중치에 대해 추가되어, 양자화 및 역양자화 과정에서의 정확도를 보장합니다.

왜 이게 좋은가?

이번 PR은 여러 측면에서 ONNX Runtime의 성능과 효율성을 크게 향상시킵니다.

1. 메모리 대역폭 및 용량 절감

가장 명확한 이점은 가중치 데이터의 메모리 사용량을 줄이는 것입니다. 4-bit에서 2-bit로 양자화하면 가중치 데이터 크기가 절반으로 줄어듭니다. 이는 다음과 같은 효과를 가져옵니다:

  • 모델 크기 감소: 모델 파일 크기가 작아져 저장 및 배포가 용이해집니다.
  • 메모리 대역폭 요구량 감소: GPU의 고정된 메모리 대역폭 내에서 더 많은 데이터를 처리할 수 있게 되어, 특히 메모리 대역폭이 병목인 연산에서 성능 향상을 기대할 수 있습니다.
  • 캐시 효율성 향상: 더 많은 가중치를 L1/L2 캐시 또는 레지스터에 올릴 수 있어, 데이터 로딩 지연 시간을 줄일 수 있습니다.

2. 연산 속도 향상 (Fused Kernel)

PR 설명에 따르면, 특히 FP16 활성화와 2-bit 가중치를 사용하는 경우, 새로운 융합(fused) 커널이 기존의 dequantize + cuBLAS 방식보다 훨씬 빠릅니다.

성능 측정 결과 (H200, FP16 activations, block_size=128 기준):

  • GEMV: 융합된 fpA_intB GEMV는 테스트된 모든 프로덕션 형태에서 기존의 수기 작성된 2-bit GEMV보다 빠르며, dequantize + cuBLAS 방식 대비 5.41-17.46 us 대 16.19-141.11 us의 성능을 보였습니다. 이는 최대 약 26배의 성능 향상입니다.
  • GEMM: 레지스터 버퍼 로테이션을 통한 최적화는 M이 {128, 512, 2048}일 때 2-bit GEMM 성능을 4.2-7.1% 향상시켰습니다.

이러한 성능 향상은 2-bit 가중치를 직접 처리하는 최적화된 CUDA 커널을 사용하고, 불필요한 역양자화 단계를 제거했기 때문에 가능합니다.

3. Compact 빌드 옵션

onnxruntime_USE_FPA_INTB_GEMM_FULL 옵션이 기본적으로 OFF 상태이며, compact 빌드 시 FP16/BF16 활성화와 INT2/INT4/INT8 가중치에 대한 scale-only GEMM/GEMV 커널을 포함하도록 변경되었습니다. 이는 별도의 설정 없이도 기본 빌드에서 2-bit 양자화의 이점을 누릴 수 있게 하여 접근성을 높입니다.

4. 일반적인 교훈

  • 양자화의 중요성: 딥러닝 모델의 성능과 효율성을 극대화하기 위해 양자화는 필수적인 기법입니다. 특히 2-bit와 같이 더 낮은 비트폭으로의 양자화는 모델 경량화의 극한을 추구하는 데 중요한 역할을 합니다.
  • 하드웨어 특화 최적화: CUDA와 같은 특정 하드웨어 아키텍처의 기능을 최대한 활용하는 커널 개발은 성능 향상의 핵심입니다. SM80/SM90 아키텍처의 특성, 메모리 레이아웃, 레지스터 활용 등을 고려한 최적화가 중요합니다.
  • 빌드 시스템과 문서의 중요성: 새로운 기능을 추가할 때는 CMake 설정과 같은 빌드 시스템을 명확하게 관리하고, 관련 문서를 최신 상태로 유지하여 사용자가 기능을 쉽게 이해하고 활용할 수 있도록 해야 합니다.
  • 정확도와 성능의 균형: 양자화는 성능 향상을 가져오지만, 정확도 손실을 최소화하는 것이 중요합니다. A = I와 같은 정확도 테스트는 양자화/역양자화 파이프라인의 무결성을 보장하는 데 필수적입니다.

리뷰 댓글 분석

리뷰 댓글들은 코드의 완성도와 테스트 커버리지를 높이는 데 중요한 역할을 했습니다.

  • 테스트 시간 및 정확도: Copilot의 지적처럼, M=512와 같이 큰 차원에서의 테스트는 전체 테스트 시간을 과도하게 소요시킬 수 있습니다. 이에 대한 해결책으로, tianleiwu가 제안한 것처럼, 프로덕션 N/K 차원을 유지하되 M을 줄이거나, 참조 구현을 최적화하여 테스트 시간을 단축하는 방안이 채택되었습니다. 이는 테스트의 실효성을 유지하면서도 효율성을 높이는 좋은 예시입니다.
  • Compact 빌드와 테스트 경로: Compact 빌드 설정에서 2-bit 연산이 올바른 경로(fpA_intB)로 실행되는지 확인하는 것이 중요했습니다. Copilot의 지적에 따라 block_size를 비트별로 적절하게 설정하고, CPU fallback을 비활성화하여 CUDA Provider의 동작을 정확히 검증하도록 수정되었습니다. tianleiwu의 수정 사항은 이러한 문제를 해결하여 CI에서 의도된 경로가 테스트되도록 보장했습니다.
  • CPU Fallback 비활성화: CUDA Provider의 동작을 정확히 테스트하기 위해 CPU fallback을 비활성화하는 것은 EP(Execution Provider)별 테스트에서 일반적인 모범 사례입니다. 이는 해당 Provider의 코드가 실제로 실행되는지 여부를 명확히 합니다.

이러한 리뷰 과정은 PR의 품질을 크게 향상시켰으며, 특히 테스트 커버리지와 정확성 측면에서 중요한 기여를 했습니다.

결론

이번 ONNX Runtime CUDA Provider의 2-bit 양자화 GEMM/GEMV 지원 추가는 모델 경량화 및 추론 성능 향상에 있어 중요한 진전입니다. 특히 LLM과 같이 MatMul 연산이 지배적인 모델에서 메모리 대역폭 병목을 완화하고 연산 속도를 크게 개선할 수 있습니다. 최적화된 커널 구현, 스마트한 빌드 옵션 관리, 그리고 철저한 테스트 및 문서화는 이러한 기술적 성과를 뒷받침합니다.

앞으로도 ONNX Runtime의 지속적인 발전과 최적화를 통해 더 많은 딥러닝 모델이 효율적으로 배포되고 활용될 수 있기를 기대합니다.

References

참고 자료

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.

댓글

관련 포스트

PR Analysis 의 다른글