DeepGEMM: Una biblioteca de código abierto con soporte eficiente para operaciones matriciales FP8 (DeepSeek Open Source Week Day 3)

Introducción general

DeepGEMM es un sistema de DeepSeek La biblioteca FP8 GEMM (Generalised Matrix Multiplication) de código abierto desarrollada por el equipo se centra en proporcionar un soporte eficiente de operaciones matriciales. Está diseñada específicamente para el núcleo tensorial de la arquitectura Hopper de NVIDIA y admite tanto operaciones matriciales comunes como operaciones GEMM agrupadas para modelos expertos mixtos (MoE). Escrita en CUDA, la librería se compila en el kernel de ejecución mediante el método de compilación ligera Just-In-Time (JIT), lo que elimina la necesidad de precompilación durante la instalación y simplifica enormemente el proceso de implantación. DeepGEMM proporciona un rendimiento extraordinario manteniendo un código limpio, lo que le permite alcanzar más de 1350 TFLOPS de capacidad de cálculo FP8 en las GPU Hopper. No sólo es adecuado para el entrenamiento de modelos de aprendizaje automático y la aceleración de la inferencia, sino que también es un recurso excelente para aprender optimización de matrices FP8 gracias a su naturaleza de código abierto y su accesibilidad.

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

 

Lista de funciones

-Admite operaciones matriciales FP8Proporciona una eficiente multiplicación de matrices generalizada FP8 (GEMM) para escenarios de computación de alto rendimiento.
-Optimización del modelo ME: Soporte para GEMM agrupado para modelos expertos híbridos, agrupados sólo para el eje M, con expertos en adaptación que comparten la misma forma de la escena.
-Compilación "justo a tiempo" (JIT)Compila el núcleo en tiempo de ejecución para adaptarlo a distintos entornos de hardware sin necesidad de precompilación.
-informática de alto rendimiento (HPC)Logró un rendimiento de cálculo FP8 de más de 1350 TFLOPS en las GPU NVIDIA Hopper.
-Diseño de código sencillo: Alrededor de 300 líneas de código central, fácil de aprender y desarrollo secundario.
-alta compatibilidadAdmite tanto GEMM normal como GEMM de paquetes con máscara, lo que resulta adecuado para diversos escenarios de inferencia.
-Código abierto y gratuitoPublicado bajo el protocolo del MIT para investigación y uso comercial.

Utilizar la ayuda

DeepGEMM es una biblioteca de aritmética matricial de código abierto diseñada para desarrolladores, principalmente para usuarios con conocimientos básicos de programación CUDA y aprendizaje automático. A continuación encontrarás una guía detallada que te ayudará a ponerte manos a la obra rápidamente y a integrarla en tu proyecto.

Proceso de instalación

DeepGEMM no requiere un complejo proceso de precompilación y puede instalarse y configurar el entorno de ejecución en unos pocos pasos:
1.Preparación medioambiental::

  • Requisitos del sistema: GPU compatibles con la arquitectura Hopper de NVIDIA (por ejemplo, H100).
  • Dependencias de software: Instale CUDA Toolkit (versión recomendada 11.8 o superior) y Python (3.8+).
  • Soporte de hardware: asegúrate de que tu dispositivo está equipado con una GPU NVIDIA con al menos 40 GB de memoria de vídeo.
    2.almacén de clones::
    Descargue el repositorio DeepGEMM localmente ejecutando el siguiente comando en un terminal:
git clone https://github.com/deepseek-ai/DeepGEMM.git**
cd DeepGEMM**
  1. Instale la dependencia:
    Utilice las herramientas de gestión de paquetes de Python para instalar las dependencias necesarias:
pip install torch numpy

DeepGEMM no requiere compilación adicional, ya que se basa en técnicas de compilación sobre la marcha y todos los núcleos se generan automáticamente en tiempo de ejecución.
4. Verifique la instalación:
Ejecute los scripts de prueba proporcionados para asegurarse de que el entorno está configurado correctamente:

python test/deep_gemm_test.py

Si la salida muestra resultados aritméticos matriciales normales, la instalación se ha realizado correctamente.

Funciones principales

1. Realización de operaciones básicas de GEMM de FP8

DeepGEMM proporciona una interfaz fácil de usar para realizar multiplicaciones matriciales FP8 no agrupadas:

  • Procedimiento operativo:
    • Importar bibliotecas y funciones:
import torch
from deep_gemm import gemm_fp8_fp8_bf16_nt
  • Prepare los datos de entrada (matrices A y B, deben estar en formato FP8):
A = torch.randn(1024, 512, dtype=torch.float8_e4m3fn).cuda()
B = torch.randn(512, 1024, dtype=torch.float8_e4m3fn).cuda()
  • Llama a la función para realizar la multiplicación de matrices:
C = gemm_fp8_fp8_bf16_nt(A, B)
print(C)
  • Advertencias:
    • La matriz de entrada debe estar ubicada en la GPU y en formato FP8 (E4M3 o E5M2).
    • La salida está en formato BF16, adecuado para cálculos posteriores o almacenamiento.

2. Agrupación de GEMM en apoyo del modelo ME

Para los usuarios que necesiten trabajar con modelos MoE, DeepGEMM proporciona soporte GEMM agrupado:

  • Procedimiento operativo:
    • Funciones GEMM de agrupación de importaciones:
from deep_gemm import m_grouped_gemm_fp8_fp8_bf16_nt_contiguous
  • Preparar los datos de entrada para el trazado continuo:
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
  • Realizar grupo GEMM:
C = m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(A, B, group_sizes)
print(C)
  • Advertencias:
    • Los ejes M de la matriz de entrada A deben agruparse por expertos y el tamaño de cada grupo debe ajustarse al tamaño del bloque M de GEMM (disponible). get_m_alignment_for_contiguous_layout() (Acceso).
    • Es necesario fijar los ejes N y K de la matriz B.

3. Agrupación de máscaras en la fase de inferencia GEMM

En la fase de decodificación de la inferencia, DeepGEMM soporta GEMM agrupado utilizando máscaras para la asignación dinámica de tokens:

  • Procedimiento operativo:
    • Importa la función de agrupación de máscaras:
from deep_gemm import m_grouped_gemm_fp8_fp8_bf16_nt_masked
  • Prepare los datos de entrada y la máscara:
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
  • Realizar agrupación de máscaras GEMM:
C = m_grouped_gemm_fp8_fp8_bf16_nt_masked(A, B, mask)
print(C)
  • Advertencias:
    • Las máscaras se utilizan para especificar qué tokens deben calcularse y son adecuadas para el razonamiento dinámico cuando los gráficos CUDA están activados.

Funciones destacadas Procedimiento de funcionamiento

Optimización y depuración de alto rendimiento

El principal punto fuerte de DeepGEMM es su eficiencia y simplicidad, que los desarrolladores pueden optimizar y depurar aún más siguiendo los pasos que se indican a continuación:

  • Ver datos de rendimiento:
    Añadir salida de registro para supervisar TFLOPS al ejecutar scripts de prueba:

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

Parámetros de ajuste*:
Ajuste el tamaño de bloque (parámetro TMA) al hardware específico para optimizar el movimiento de datos y el solapamiento computacional, consulte el ejemplo en la carpeta test/ de la documentación.
Aprendizaje y expansión*:
El código central se encuentra en deep_gemm/gemm_kernel.cu, unas 300 líneas, los desarrolladores pueden leerlo directamente y modificarlo para adaptarlo a sus necesidades.
Recomendaciones de uso
requisitos de hardware*: Actualmente sólo es compatible con GPUs de arquitectura NVIDIA Hopper, otras arquitecturas aún no están adaptadas.
referencia documental*: Puede encontrar descripciones detalladas de las funciones y código de ejemplo en las carpetas README.md y test/ de los repositorios de GitHub.
Apoyo comunitario*: Si encuentra problemas, envíe sus comentarios a la página de problemas de GitHub y el equipo de DeepSeek responderá positivamente.
Con estos pasos, podrás integrar fácilmente DeepGEMM en tus proyectos de aprendizaje automático y disfrutar de su eficiente potencia de cálculo matricial FP8.

© declaración de copyright

Artículos relacionados

Sin comentarios

Debe iniciar sesión para participar en los comentarios.
Acceder ahora
ninguno
Sin comentarios...