DeepGEMMは、AIトレーニングおよび推論のためのFP8 GEMMライブラリで、特に密行列および混合エキスパート(MoE)行列演算用に設計されており、DeepSeek V3およびR1モデルのトレーニングと推論を強力にサポートします。
DeepSeekは相次いでオープンソース化している。 フラッシュMLA そして ディープEPそして本日、次のように発表した。 ディープゲーエムHopperアーキテクチャGPU向けに最適化された行列乗算ライブラリです。このライブラリは、標準的な行列計算とMoE(mixture-of-experts)計算をサポートし、DeepSeek-V3/R1のトレーニングと推論を強力にサポートします。
ディープゲーエム deepGEMM はシンプルかつ効率的に設計されており、コアコードはわずか 300 行程度です。このライブラリは3つのデータアラインメントをサポートしています。標準的なアラインメントと、ハイブリッドエキスパートモデル用に設計された2つの特別なアラインメント(シーケンシャルとマスク)です。deepGEMMはオンザフライコンパイルを採用しているため、インストール時にコンパイルする必要がなく、明快で理解しやすいコード構造になっているため、GPU最適化技術の学習に最適です。
DeepGEMMは、さまざまな計算シナリオで優れた性能を発揮します。標準的な行列の乗算では、CUTLASS 3.6に基づく最適化された実装と比較して、1.0倍から2.7倍のスピードアップが得られました。最も大幅な高速化(最大2.7倍)は、小バッチのデータ(M=64または128)に対して達成された。ハイブリッドエキスパートモデルの計算では、DeepGEMMが提供する2つの特別なデータアラインメントも大きな利点となります。シーケンシャル配列は、学習とバッチ推論の両方のフェーズに適しており、約1.1~1.2倍の高速化を実現します。一方、マスク配列は、リアルタイム推論用に設計されており、CUDAグラフ技術の使用をサポートし、同じく1.1~1.2倍の高速化を実現します。
M | N | K | 計算 | メモリ帯域幅 | スピードアップ |
---|---|---|---|---|---|
64 | 2112 | 7168 | 206 TFLOPS | 1688GB/秒 | 2.7x |
64 | 24576 | 1536 | 289 TFLOPS | 2455 GB/秒 | 1.7x |
64 | 32768 | 512 | 219 TFLOPS | 2143 GB/秒 | 1.8x |
64 | 7168 | 16384 | 336 TFLOPS | 2668 GB/秒 | 1.4x |
64 | 4096 | 7168 | 287 TFLOPS | 2320 GB/秒 | 1.4x |
64 | 7168 | 2048 | 295 TFLOPS | 2470GB/秒 | 1.7x |
128 | 2112 | 7168 | 352 TFLOPS | 1509 GB/秒 | 2.4x |
128 | 24576 | 1536 | 535 TFLOPS | 2448GB/秒 | 1.6x |
128 | 32768 | 512 | 358 TFLOPS | 2103GB/秒 | 1.5x |
128 | 7168 | 16384 | 645 TFLOPS | 2604GB/秒 | 1.4x |
128 | 4096 | 7168 | 533 TFLOPS | 2221GB/秒 | 2.0x |
128 | 7168 | 2048 | 510 TFLOPS | 2277GB/秒 | 1.7x |
4096 | 2112 | 7168 | 1058 TFLOPS | 527GB/秒 | 1.1x |
4096 | 24576 | 1536 | 990 TFLOPS | 786 GB/秒 | 1.0x |
4096 | 32768 | 512 | 590 TFLOPS | 1232GB/秒 | 1.0x |
4096 | 7168 | 16384 | 1358 TFLOPS | 343GB/秒 | 1.2x |
4096 | 4096 | 7168 | 1304 TFLOPS | 500GB/秒 | 1.1x |
4096 | 7168 | 2048 | 1025 TFLOPS | 697 GB/秒 | 1.1x |
#Gグループ | M/グループ | N | K | 計算 | メモリ帯域幅 | スピードアップ |
---|---|---|---|---|---|---|
4 | 8192 | 4096 | 7168 | 1297 TFLOPS | 418 GB/秒 | 1.2x |
4 | 8192 | 7168 | 2048 | 1099 TFLOPS | 681 GB/秒 | 1.2x |
8 | 4096 | 4096 | 7168 | 1288 TFLOPS | 494 GB/秒 | 1.2x |
8 | 4096 | 7168 | 2048 | 1093 TFLOPS | 743GB/秒 | 1.1x |
#Gグループ | M/グループ | N | K | 計算 | メモリ帯域幅 | スピードアップ |
---|---|---|---|---|---|---|
1 | 1024 | 4096 | 7168 | 1233 TFLOPS | 924 GB/秒 | 1.2x |
1 | 1024 | 7168 | 2048 | 925 TFLOPS | 968 GB/秒 | 1.2x |
2 | 512 | 4096 | 7168 | 1040 TFLOPS | 1288 GB/秒 | 1.2x |
2 | 512 | 7168 | 2048 | 916 TFLOPS | 1405 GB/秒 | 1.2x |
4 | 256 | 4096 | 7168 | 932 TFLOPS | 2064GB/秒 | 1.1x |
4 | 256 | 7168 | 2048 | 815 TFLOPS | 2047 GB/秒 | 1.2x |
DeepGEMMを使用するには、sm_90aをサポートするHopperアーキテクチャGPU、Python 3.8以上、CUDA 12.3以上(最高のパフォーマンスを得るには12.8以上を推奨)、PyTorch 2.1以上、CUTLASS 3.6以上が必要です。
開発
#サブモジュールはクローンでなければならない
git clone --recursive [email protected]:deepseek-ai/DeepGEMM.git
# サードパーティ (CUTLASS および CuTe) のインクルードディレクトリにシンボリックリンクを作成する。
python setup.py develop
# テストJITコンパイル
python tests/test_jit.py
# すべての GEMM インプリメント(ノーマル、連続グループ化、マスクグループ化)をテストする。
python tests/test_core.py
インストール
python setup.py インストール
最後に ディープ・ジェム それで終わりだ!