DeepGEMM : Une bibliothèque Open Source avec un support efficace pour les opérations matricielles du FP8 (DeepSeek Open Source Week Day 3)

Introduction générale

DeepGEMM est un système de DeepSeek La bibliothèque FP8 GEMM (Generalised Matrix Multiplication) open source développée par l'équipe est axée sur la prise en charge efficace des opérations matricielles. Elle est spécialement conçue pour le Tensor Core de l'architecture NVIDIA Hopper et prend en charge les opérations matricielles courantes et les opérations GEMM groupées pour les modèles experts mixtes (MoE). Écrite en CUDA, la bibliothèque est compilée dans le noyau d'exécution à l'aide d'une compilation JIT (Just-In-Time) légère, ce qui élimine le besoin de précompilation à l'installation et simplifie grandement le processus de déploiement. DeepGEMM offre des performances exceptionnelles tout en conservant un code propre, atteignant plus de 1 350 TFLOPS de puissance de calcul en FP8 sur les GPU Hopper. DeepGEMM est non seulement adapté à l'entraînement de modèles d'apprentissage automatique et à l'accélération de l'inférence, mais il constitue également une excellente ressource pour l'apprentissage de l'optimisation de matrices FP8 en raison de sa nature open source et de son accessibilité.

DeepGEMM:高效支持FP8矩阵运算的开源库(DeepSeek 开源周第三天)

 

Liste des fonctions

-Prise en charge des opérations matricielles FP8Le système d'information sur les matrices (GEMM) est un outil efficace pour la multiplication de matrices généralisée (FP8) dans des scénarios de calcul à haute performance.
-Optimisation du modèle MoEPrise en charge des GEMM groupés pour les modèles experts hybrides, groupés uniquement pour l'axe M, avec des experts en adaptation partageant la même forme de la scène.
-Compilation juste à temps (JIT)Compiler le noyau au moment de l'exécution pour l'adapter à différents environnements matériels sans précompilation.
-le calcul à haute performance (HPC)Le groupe a obtenu un débit de calcul de plus de 1350 TFLOPS sur les GPU NVIDIA Hopper pour le FP8.
-Conception d'un code simpleLe système de gestion de l'information : Environ 300 lignes de code de base, faciles à apprendre et à développer en second lieu.
-haute compatibilitéIl prend en charge à la fois le GEMM normal et le GEMM par paquets avec masque, ce qui convient à divers scénarios d'inférence.
-Open source et gratuit: Publié sous le protocole du MIT pour la recherche et l'utilisation commerciale.

Utiliser l'aide

DeepGEMM est une bibliothèque d'arithmétique matricielle open source conçue pour les développeurs, principalement pour les utilisateurs ayant des connaissances de base en programmation CUDA et en apprentissage automatique. Ce qui suit est un guide détaillé pour vous aider à démarrer rapidement et à l'intégrer dans votre projet.

Processus d'installation

DeepGEMM ne nécessite pas de processus de précompilation complexe et peut être installé et l'environnement d'exécution configuré en quelques étapes seulement :
1.Préparation de l'environnement: :

  • Configuration requise : GPU compatibles avec l'architecture NVIDIA Hopper (par exemple H100).
  • Dépendances logicielles : Installer CUDA Toolkit (version recommandée 11.8 ou supérieure) et Python (3.8+).
  • Support matériel : Assurez-vous que votre appareil est équipé d'un GPU NVIDIA avec au moins 40 Go de mémoire vidéo.
    2.entrepôt de clones: :
    Téléchargez le dépôt DeepGEMM localement en exécutant la commande suivante dans un terminal :
git clone https://github.com/deepseek-ai/DeepGEMM.git**
cd DeepGEMM**
  1. Installer la dépendance :
    Utilisez les outils de gestion des paquets de Python pour installer les dépendances nécessaires :
pip install torch numpy

DeepGEMM lui-même ne nécessite aucune compilation supplémentaire, car il s'appuie sur des techniques de compilation à la volée et tous les noyaux sont automatiquement générés au moment de l'exécution.
4. vérifier l'installation :
Exécutez les scripts de test fournis pour vous assurer que l'environnement est correctement configuré :

python test/deep_gemm_test.py

Si la sortie montre des résultats arithmétiques matriciels normaux, l'installation a réussi.

Principales fonctions

1. effectuer les opérations de base du GEMM FP8

DeepGEMM fournit une interface facile à utiliser pour effectuer des multiplications matricielles FP8 non groupées :

  • Procédure d'exploitation :
    • Importer des bibliothèques et des fonctions :
import torch
from deep_gemm import gemm_fp8_fp8_bf16_nt
  • Préparer les données d'entrée (les matrices A et B doivent être au format FP8) :
A = torch.randn(1024, 512, dtype=torch.float8_e4m3fn).cuda()
B = torch.randn(512, 1024, dtype=torch.float8_e4m3fn).cuda()
  • Appeler la fonction pour effectuer la multiplication de la matrice :
C = gemm_fp8_fp8_bf16_nt(A, B)
print(C)
  • Mises en garde :
    • La matrice d'entrée doit être située sur le GPU et au format FP8 (E4M3 ou E5M2).
    • La sortie est au format BF16, adapté aux calculs ultérieurs ou au stockage.

2. le regroupement du GEMM pour soutenir le modèle des ministères de l'environnement

Pour les utilisateurs qui doivent travailler avec des modèles MoE, DeepGEMM fournit un support GEMM groupé :

  • Procédure d'exploitation :
    • Importation des fonctions de regroupement du GEMM :
from deep_gemm import m_grouped_gemm_fp8_fp8_bf16_nt_contiguous
  • Préparer les données d'entrée pour la mise en page en continu :
A = torch.randn(4096, 512, dtype=torch.float8_e4m3fn).cuda()  # 多个专家的输入拼接
B = torch.randn(512, 1024, dtype=torch.float8_e4m3fn).cuda()
group_sizes = [1024, 1024, 1024, 1024]  # 每个专家的 token
  • Exécuter le groupe GEMM :
C = m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(A, B, group_sizes)
print(C)
  • Mises en garde :
    • Les axes M de la matrice d'entrée A doivent être regroupés par des experts et la taille de chaque groupe doit être alignée sur la taille du bloc M du GEMM (disponible). get_m_alignment_for_contiguous_layout() (Accès).
    • Les axes N et K de la matrice B doivent être fixés.

3. regroupement des masques au stade de l'inférence GEMM

Dans la phase de décodage de l'inférence, DeepGEMM prend en charge les GEMM groupés en utilisant des masques pour l'allocation dynamique des jetons :

  • Procédure d'exploitation :
    • Importer la fonction de regroupement des masques :
from deep_gemm import m_grouped_gemm_fp8_fp8_bf16_nt_masked
  • Préparer les données d'entrée et le masque :
A = torch.randn(4096, 512, dtype=torch.float8_e4m3fn).cuda()
B = torch.randn(512, 1024, dtype=torch.float8_e4m3fn).cuda()
mask = torch.ones(4096, dtype=torch.bool).cuda()  # 掩码指示有效 token
  • Effectuer le regroupement des masques GEMM :
C = m_grouped_gemm_fp8_fp8_bf16_nt_masked(A, B, mask)
print(C)
  • Mises en garde :
    • Les masques sont utilisés pour spécifier quels jetons doivent être calculés et conviennent au raisonnement dynamique lorsque les graphes CUDA sont activés.

Fonctions vedettes Procédure d'utilisation

Optimisation et débogage de haute performance

La force principale de DeepGEMM est son efficacité et sa simplicité, que les développeurs peuvent optimiser et déboguer en suivant les étapes ci-dessous :

  • Visualiser les données de performance :
    Ajout d'une sortie de journal pour surveiller les TFLOPS lors de l'exécution des scripts de test :

    import logging
    logging.basicConfig(level=logging.INFO)
    C = gemm_fp8_fp8_bf16_nt(A, B)
    

Paramètres de réglage* :
Ajustez la taille du bloc (paramètre TMA) en fonction du matériel spécifique pour optimiser le mouvement des données et le chevauchement des calculs, voir l'exemple dans le dossier test/ de la documentation.
Apprentissage et expansion* :
Le code principal se trouve dans deep_gemm/gemm_kernel.cu, environ 300 lignes, les développeurs peuvent le lire directement et le modifier pour l'adapter à leurs besoins.
Recommandations d'utilisation
exigences en matière de matériel* : Actuellement, seuls les GPU NVIDIA Hopper sont pris en charge, les autres architectures ne sont pas encore adaptées.
référence de la documentation* : Des descriptions détaillées des fonctions et des exemples de code peuvent être trouvés dans les dossiers README.md et test/ des dépôts GitHub.
Soutien communautaire* : Si vous rencontrez des problèmes, vous pouvez soumettre des commentaires sur la page GitHub Issues et l'équipe DeepSeek y répondra positivement.
Avec ces étapes, vous pouvez facilement intégrer DeepGEMM dans vos projets d'apprentissage automatique et profiter de sa puissance de calcul matricielle FP8.

© déclaration de droits d'auteur

Articles connexes

Pas de commentaires

Vous devez être connecté pour participer aux commentaires !
S'inscrire maintenant
aucun
Pas de commentaires...