ディープゲーエム

AIのトレーニングと推論

DeepGEMMは、AIトレーニングおよび推論のためのFP8 GEMMライブラリで、特に密行列および混合エキスパート(MoE)行列演算用に設計されており、DeepSeek V3およびR1モデルのトレーニングと推論を強力にサポートします。

ディープゲーエム

DeepGEMMとは?

DeepSeekは相次いでオープンソース化している。 フラッシュMLA そして ディープEPそして本日、次のように発表した。 ディープゲーエムHopperアーキテクチャGPU向けに最適化された行列乗算ライブラリです。このライブラリは、標準的な行列計算とMoE(mixture-of-experts)計算をサポートし、DeepSeek-V3/R1のトレーニングと推論を強力にサポートします。

ディープゲーエム deepGEMM はシンプルかつ効率的に設計されており、コアコードはわずか 300 行程度です。このライブラリは3つのデータアラインメントをサポートしています。標準的なアラインメントと、ハイブリッドエキスパートモデル用に設計された2つの特別なアラインメント(シーケンシャルとマスク)です。deepGEMMはオンザフライコンパイルを採用しているため、インストール時にコンパイルする必要がなく、明快で理解しやすいコード構造になっているため、GPU最適化技術の学習に最適です。

 

DeepGEMMのパフォーマンス

DeepGEMMは、さまざまな計算シナリオで優れた性能を発揮します。標準的な行列の乗算では、CUTLASS 3.6に基づく最適化された実装と比較して、1.0倍から2.7倍のスピードアップが得られました。最も大幅な高速化(最大2.7倍)は、小バッチのデータ(M=64または128)に対して達成された。ハイブリッドエキスパートモデルの計算では、DeepGEMMが提供する2つの特別なデータアラインメントも大きな利点となります。シーケンシャル配列は、学習とバッチ推論の両方のフェーズに適しており、約1.1~1.2倍の高速化を実現します。一方、マスク配列は、リアルタイム推論用に設計されており、CUDAグラフ技術の使用をサポートし、同じく1.1~1.2倍の高速化を実現します。

密なモデルに対する通常のGEMM

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

MoEモデルのグループ化されたGEMM(連続レイアウト)

#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

MoEモデルのグループ化GEMM(マスクレイアウト)

#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の使い方は?

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 インストール

最後に ディープ・ジェム それで終わりだ!