DeepGEMM

AI-Training und Inferenz

DeepGEMM ist eine FP8 GEMM-Bibliothek für KI-Training und Inferenz, die speziell für dichte und Mixture of Experts (MoE)-Matrixoperationen entwickelt wurde und robuste Unterstützung für das Training und die Inferenz von DeepSeek V3- und R1-Modellen bietet.

DeepGEMM

Was ist DeepGEMM?

DeepSeek hat sukzessive das Open-Sourcing FlashMLA und DeepEPund heute stellt sie vor DeepGEMMeine Bibliothek für Matrixmultiplikation, die speziell für GPUs mit Hopper-Architektur optimiert wurde. Diese Bibliothek unterstützt Standard-Matrixberechnungen und Mixture-of-Experts (MoE)-Berechnungen und bietet robuste Unterstützung für das Training und die Inferenz von DeepSeek-V3/R1 und erreicht eine hohe Leistung von 1350+ FP8 TFLOPS auf Hopper-GPUs.

DeepGEMM ist so konzipiert, dass sie einfach und effizient ist, mit nur etwa 300 Zeilen Kerncode, und dabei bestehende Lösungen in den meisten Matrixgrößen übertrifft. Die Bibliothek unterstützt drei Datenausrichtungen: eine Standardausrichtung und zwei spezielle Ausrichtungen (sequentiell und maskiert), die für hybride Expertenmodelle entwickelt wurden. deepGEMM verwendet eine On-the-fly-Kompilierung, wodurch eine Kompilierung bei der Installation überflüssig wird, und verfügt über eine klare, leicht verständliche Codestruktur, die sich ideal für das Erlernen von GPU-Optimierungstechniken eignet.

 

DeepGEMM-Leistung

DeepGEMM schneidet in einer Vielzahl von Berechnungsszenarien gut ab. Bei der Standard-Matrixmultiplikation liegen die Geschwindigkeitssteigerungen zwischen dem 1,0- und 2,7-fachen im Vergleich zur optimierten Implementierung auf der Grundlage von CUTLASS 3.6. Die größten Geschwindigkeitssteigerungen, bis zum 2,7-fachen, wurden bei kleinen Datenmengen (M=64 oder 128) erzielt. Für die Berechnung von hybriden Expertenmodellen bieten die beiden speziellen Datenanordnungen von DeepGEMM ebenfalls deutliche Vorteile. Die sequentielle Anordnung eignet sich sowohl für die Trainings- als auch für die Batch-Inferenzphase, wobei Geschwindigkeitssteigerungen von etwa dem 1,1- bis 1,2-fachen erreicht werden, während die maskierte Anordnung für die Echtzeit-Inferenz konzipiert ist und die Verwendung von CUDA-Graphtechniken unterstützt, ebenfalls mit Geschwindigkeitssteigerungen von 1,1- bis 1,2-fach.

Normale GEMMs für dichte Modelle

M N K Berechnung Speicher-Bandbreite Beschleunigung
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

Gruppierte GEMMs für MoE-Modelle (zusammenhängendes Layout)

#Gruppen M pro Gruppe N K Berechnung Speicher-Bandbreite Beschleunigung
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

Gruppierte GEMMs für MoE-Modelle (maskiertes Layout)

#Gruppen M pro Gruppe N K Berechnung Speicher-Bandbreite Beschleunigung
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

Wie verwendet man DeepGEMM?

Um DeepGEMM zu verwenden, benötigen Sie Grafikprozessoren der Hopper-Architektur mit sm_90a-Unterstützung, Python 3.8 oder höher, CUDA 12.3 oder höher (12.8 oder höher wird für beste Leistung empfohlen), PyTorch 2.1 oder höher und CUTLASS 3.6 oder höher.

Entwicklung

# Submodul muss geklont werden
git clone --recursive [email protected]:deepseek-ai/DeepGEMM.git

# Symbolische Links für Include-Verzeichnisse von Drittanbietern (CUTLASS und CuTe) erstellen
python setup.py entwickeln

# Test JIT-Kompilierung
python tests/test_jit.py

# Test aller GEMM-Implementierungen (normal, zusammenhängend gruppiert und maskiert gruppiert)
python tests/test_core.py

Einrichtung

python setup.py installieren

Schließlich importieren Sie tief_gemm und Sie sind fertig!