DeepGEMM

Formação e inferência de IA

DeepGEMM é uma biblioteca FP8 GEMM para treinamento e inferência de IA, projetada especificamente para operações de matriz densa e mistura de especialistas (MoE), fornecendo suporte robusto para o treinamento e inferência de modelos DeepSeek V3 e R1.

DeepGEMM

O que é a DeepGEMM?

O DeepSeek tem vindo a abrir sucessivamente FlashMLA e DeepEPe hoje apresenta DeepGEMM, uma biblioteca de multiplicação de matrizes optimizada especificamente para GPUs de arquitetura Hopper. Esta biblioteca suporta computação de matriz padrão e computação de mistura de especialistas (MoE), fornecendo suporte robusto para o treinamento e inferência do DeepSeek-V3/R1, alcançando alto desempenho de 1350+ FP8 TFLOPS em GPUs Hopper.

DeepGEMM A biblioteca deepGEMM foi projetada para ser simples e eficiente, com apenas cerca de 300 linhas de código principal, enquanto supera as soluções existentes na maioria dos tamanhos de matriz. A biblioteca suporta três alinhamentos de dados: um alinhamento padrão e dois alinhamentos especiais (sequencial e mascarado) concebidos para modelos especializados híbridos. deepGEMM usa compilação on-the-fly, eliminando a necessidade de compilar no momento da instalação, e tem uma estrutura de código clara e fácil de entender que o torna ideal para aprender técnicas de otimização de GPU.

 

Desempenho da DeepGEMM

O DeepGEMM tem um bom desempenho numa variedade de cenários computacionais. Para a multiplicação de matrizes padrão, os aumentos de velocidade variam de 1,0 a 2,7 vezes em comparação com a implementação optimizada baseada no CUTLASS 3.6. Os aumentos de velocidade mais significativos, até 2,7 vezes, foram obtidos para pequenos lotes de dados (M=64 ou 128). Para a computação de modelos especializados híbridos, os dois alinhamentos de dados especiais oferecidos pelo DeepGEMM também oferecem vantagens significativas. O arranjo sequencial é adequado para as fases de treinamento e inferência em lote, com aumentos de velocidade de cerca de 1,1 a 1,2 vezes, enquanto o arranjo mascarado é projetado para inferência em tempo real e suporta o uso de técnicas gráficas CUDA, também com aumentos de velocidade de 1,1 a 1,2 vezes.

GEMMs normais para modelos densos

M N K Computação Largura de banda da memória Aceleração
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

GEMMs agrupadas para modelos MoE (layout contíguo)

#Grupos M por grupo N K Computação Largura de banda da memória Aceleração
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 agrupadas para modelos MoE (layout mascarado)

#Grupos M por grupo N K Computação Largura de banda da memória Aceleração
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

Como utilizar o DeepGEMM?

Para usar o DeepGEMM, são necessárias GPUs de arquitetura Hopper com suporte a sm_90a, Python 3.8 ou superior, CUDA 12.3 ou superior (recomenda-se 12.8 ou superior para melhor desempenho), PyTorch 2.1 ou superior e CUTLASS 3.6 ou superior.

Desenvolvimento

# O submódulo deve ser clonado
git clone --recursive [email protected]:deepseek-ai/DeepGEMM.git

# Criar ligações simbólicas para diretórios de inclusão de terceiros (CUTLASS e CuTe)
python setup.py desenvolver

# Teste de compilação JIT
python tests/test_jit.py

# Testar todas as aplicações GEMM (normal, agrupada por contiguidade e agrupada por máscara)
python tests/test_core.py

Instalação

python setup.py install

Finalmente, importar deep_gemm e já está!