DeepGEMM

Formation et inférence en IA

DeepGEMM est une bibliothèque GEMM FP8 pour l'entraînement et l'inférence d'IA, conçue spécifiquement pour les opérations matricielles denses et de mélange d'experts (MoE), fournissant un support robuste pour l'entraînement et l'inférence des modèles DeepSeek V3 et R1.

DeepGEMM

Qu'est-ce que DeepGEMM ?

DeepSeek a successivement mis en libre accès les éléments suivants FlashMLA et Profondeur de l'EPElle présente aujourd'hui DeepGEMMLa bibliothèque de multiplication matricielle de DeepSeek-V3/R1 est optimisée spécifiquement pour les GPU de l'architecture Hopper. Cette bibliothèque prend en charge le calcul matriciel standard et le calcul du mélange d'experts (MoE), fournissant un support solide pour l'entraînement et l'inférence de DeepSeek-V3/R1, atteignant des performances élevées de 1350+ FP8 TFLOPS sur les GPU Hopper.

DeepGEMM est conçue pour être simple et efficace, avec seulement 300 lignes de code, tout en surpassant les solutions existantes dans la plupart des tailles de matrices. La bibliothèque supporte trois alignements de données : un alignement standard et deux alignements spéciaux (séquentiel et masqué) conçus pour les modèles experts hybrides. deepGEMM utilise la compilation à la volée, éliminant le besoin de compiler au moment de l'installation, et possède une structure de code claire et facile à comprendre qui la rend idéale pour l'apprentissage des techniques d'optimisation GPU.

 

Performance de DeepGEMM

DeepGEMM donne de bons résultats dans divers scénarios de calcul. Pour la multiplication matricielle standard, les gains de vitesse vont de 1,0 à 2,7 fois par rapport à l'implémentation optimisée basée sur CUTLASS 3.6. Les accélérations les plus significatives, jusqu'à 2,7 fois, ont été obtenues pour de petits lots de données (M=64 ou 128). Pour le calcul des modèles experts hybrides, les deux alignements de données spéciaux proposés par DeepGEMM offrent également des avantages significatifs. L'arrangement séquentiel convient aux phases d'entraînement et d'inférence par lots, avec des accélérations d'environ 1,1 à 1,2 fois, tandis que l'arrangement masqué est conçu pour l'inférence en temps réel et prend en charge l'utilisation des techniques graphiques CUDA, également avec des accélérations de 1,1 à 1,2 fois.

GEMMs normaux pour les modèles denses

M N K Calcul Largeur de bande de la mémoire Accélération
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 Go/s 1.1x
4096 7168 2048 1025 TFLOPS 697 GB/s 1.1x

GEMMs groupés pour les modèles MoE (disposition contiguë)

#Groupes M par groupe N K Calcul Largeur de bande de la mémoire Accélération
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

GEMMs groupés pour les modèles MoE (disposition masquée)

#Groupes M par groupe N K Calcul Largeur de bande de la mémoire Accélération
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

Comment utiliser DeepGEMM ?

Pour utiliser DeepGEMM, vous avez besoin de GPU à architecture Hopper avec support sm_90a, Python 3.8 ou supérieur, CUDA 12.3 ou supérieur (12.8 ou supérieur est recommandé pour de meilleures performances), PyTorch 2.1 ou supérieur, et CUTLASS 3.6 ou supérieur.

Développement

# Le sous-module doit être cloné
git clone --recursive [email protected]:deepseek-ai/DeepGEMM.git

# Créer des liens symboliques pour les répertoires d'inclusion des tiers (CUTLASS et CuTe)
python setup.py develop

# Test Compilation JIT
python tests/test_jit.py

# Tester toutes les implémentations du GEMM (normal, groupé par contiguïté et groupé par masquage)
python tests/test_core.py

Installation

python setup.py install

Enfin, l'importation deep_gemm et le tour est joué !