Hugging Face가 pytorch_block_sparse 라이브러리를 개발해 희소 행렬 기반 선형 계층으로 언어 모델을 75% 희소도에서 4배 메모리 감소 및 2배 속도 향상 달성

Context

기존 PyTorch 희소 행렬 구현은 조밀 행렬 대비 약 10배 느린 성능을 보여주고 있었다. 또한 공식 PyTorch 희소 대수 연산 지원이 부족하여 신경망 프루닝을 통한 효율성 개선이 실제로는 활용되기 어려웠다.

Technical Solution

• BlockSparseLinear 모듈 제공: torch.nn.Linear의 드롭인 대체 모듈로 기존 코드 변경 최소화
• CUTLASS 기반 CUDA 커널 구현: NVIDIA의 CUTLASS를 사용한 C++ CUDA 템플릿으로 블록 희소 행렬 곱셈 구현
• BlockSparseModelPatcher 도구 제공: 기존 모델을 소스 코드 수정 없이 실시간으로 희소화 가능
• Ampere Tensor Core 호환성 설계: 블록 희소성이 100% Tensor Core 요구사항과 호환되도록 구현
• 디스틸레이션 및 양자화와의 조합 지원: 추가 최적화 기법들과 함께 사용 가능

Impact

• 희소 행렬 성능: cuBLAS 최적화 조밀 행렬 대비 약 2배 느림 (PyTorch 기본 희소 행렬의 1/10 속도)
• 75% 희소도에서 메모리 사용량 4배 감소
• 75% 희소도에서 처리 속도 약 2배 향상
• PyTorch 기본 희소 행렬 구현 대비 1배 이상 성능 향상

Key Takeaway

커스텀 CUDA 커널 기반 희소 행렬 구현을 통해 프레임워크 수준의 제약을 우회하면, 드롭인 인터페이스로 기존 모델에 투명하게 적용 가능한 효율성 개선을 달성할 수 있다. 프루닝된 패턴이 고정적이지만 향후 학습 중 최적화 기능이 추가되면 성능 향상 폭이 더욱 커질 것으로 예상된다.