DeepGEMM

AI 학습 및 추론

DeepGEMM은 AI 훈련 및 추론을 위한 FP8 GEMM 라이브러리로, 밀도 및 전문가 혼합(MoE) 매트릭스 작업을 위해 특별히 설계되어 DeepSeek V3 및 R1 모델의 훈련 및 추론을 강력하게 지원합니다.

DeepGEMM

DeepGEMM이란 무엇인가요?

딥시크는 연속적으로 오픈소스를 공개해 왔습니다. FlashMLA 그리고 DeepEP그리고 오늘은 DeepGEMM는 Hopper 아키텍처 GPU에 특별히 최적화된 행렬 곱셈 라이브러리입니다. 이 라이브러리는 표준 행렬 연산과 전문가 혼합(MoE) 연산을 지원하여 DeepSeek-V3/R1의 훈련 및 추론을 강력하게 지원하며, Hopper GPU에서 1350+ FP8 TFLOPS의 고성능을 달성합니다.

DeepGEMM 은 약 300줄의 핵심 코드만으로 간단하고 효율적으로 설계되었으며, 대부분의 매트릭스 크기에서 기존 솔루션보다 뛰어난 성능을 발휘합니다. 라이브러리는 표준 정렬과 하이브리드 전문가 모델을 위해 설계된 두 가지 특수 정렬(순차 및 마스크)의 세 가지 데이터 정렬을 지원합니다. deepGEMM은 즉시 컴파일을 사용하므로 설치 시 컴파일할 필요가 없으며, 명확하고 이해하기 쉬운 코드 구조를 갖추고 있어 GPU 최적화 기술을 학습하는 데 이상적입니다.

 

DeepGEMM 성능

DeepGEMM은 다양한 계산 시나리오에서 우수한 성능을 발휘합니다. 표준 행렬 곱셈의 경우, CUTLASS 3.6에 기반한 최적화된 구현에 비해 속도가 1.0~2.7배 빨라졌습니다. 가장 큰 속도 향상은 소량의 데이터 배치(M=64 또는 128)에서 최대 2.7배까지 달성되었습니다. 하이브리드 전문가 모델 계산의 경우, DeepGEMM이 제공하는 두 가지 특수 데이터 정렬도 상당한 이점을 제공합니다. 순차 배열은 훈련 및 배치 추론 단계 모두에 적합하며 약 1.1~1.2배의 속도 향상을 제공하며, 마스크 배열은 실시간 추론을 위해 설계되었으며 CUDA 그래프 기법의 사용을 지원하며 역시 1.1~1.2배의 속도 향상을 제공합니다.

고밀도 모델용 일반 GEMM

M N K 계산 메모리 대역폭 속도 향상
64 2112 7168 206 TFLOPS 1688 GB/s 2.7x
64 24576 1536 289 TFLOPS 2455 GB/s 1.7x
64 32768 512 219 TFLOPS 2143 GB/s 1.8x
64 7168 16384 336 TFLOPS 2668 GB/s 1.4x
64 4096 7168 287 TFLOPS 2320 GB/s 1.4x
64 7168 2048 295 TFLOPS 2470 GB/s 1.7x
128 2112 7168 352 TFLOPS 1509 GB/s 2.4x
128 24576 1536 535 TFLOPS 2448 GB/s 1.6x
128 32768 512 358 TFLOPS 2103 GB/s 1.5x
128 7168 16384 645 TFLOPS 2604 GB/s 1.4x
128 4096 7168 533 TFLOPS 2221 GB/s 2.0x
128 7168 2048 510 TFLOPS 2277 GB/s 1.7x
4096 2112 7168 1058 TFLOPS 527 GB/s 1.1x
4096 24576 1536 990 TFLOPS 786 GB/s 1.0x
4096 32768 512 590 TFLOPS 1232 GB/s 1.0x
4096 7168 16384 1358 TFLOPS 343 GB/s 1.2x
4096 4096 7168 1304 TFLOPS 500 GB/s 1.1x
4096 7168 2048 1025 TFLOPS 697 GB/s 1.1x

MoE 모델용 그룹화된 GEMM(연속 레이아웃)

#그룹 그룹당 M N K 계산 메모리 대역폭 속도 향상
4 8192 4096 7168 1297 TFLOPS 418 GB/s 1.2x
4 8192 7168 2048 1099 TFLOPS 681 GB/s 1.2x
8 4096 4096 7168 1288 TFLOPS 494 GB/s 1.2x
8 4096 7168 2048 1093 TFLOPS 743 GB/s 1.1x

MoE 모델용 그룹화된 GEMM(마스킹 레이아웃)

#그룹 그룹당 M N K 계산 메모리 대역폭 속도 향상
1 1024 4096 7168 1233 TFLOPS 924 GB/s 1.2x
1 1024 7168 2048 925 TFLOPS 968 GB/s 1.2x
2 512 4096 7168 1040 TFLOPS 1288 GB/s 1.2x
2 512 7168 2048 916 TFLOPS 1405 GB/s 1.2x
4 256 4096 7168 932 TFLOPS 2064 GB/s 1.1x
4 256 7168 2048 815 TFLOPS 2047 GB/s 1.2x

DeepGEMM은 어떻게 사용하나요?

DeepGEMM을 사용하려면 sm_90a를 지원하는 Hopper 아키텍처 GPU, Python 3.8 이상, CUDA 12.3 이상(최상의 성능을 위해 12.8 이상 권장), PyTorch 2.1 이상, CUTLASS 3.6 이상이 필요합니다.

개발

# 서브모듈을 복제해야 합니다.
git clone --recursive [email protected]:deepseek-ai/DeepGEMM.git

# 타사(컷클래스 및 쿠테)에 대한 심볼릭 링크에 디렉터리를 포함하도록 합니다.
파이썬 setup.py 개발

# 테스트 JIT 컴파일
파이썬 테스트/test_jit.py

# 모든 GEMM 구현(일반, 인접 그룹 및 마스크 그룹)을 테스트합니다.
파이썬 테스트/test_core.py

설치

파이썬 setup.py 설치

마지막으로 가져오기 deep_gemm 완료되었습니다!