ここ9ヶ月間私たちが続けてきたのはこうです:MoEトレーニングは~2倍速く、メモリも~2倍減らす!ハイライト: - MoEは通常、現代モデルで最も多くの時間とメモリを消費します。実は、MoEのバックワードパスを数学的に書き換えれば、前方に格納すべき活性化メモリを~2倍減らすことで、追加のマトムル再計算なしで同じ勾配が得られます。この結果はアルゴリズムとシステムの洞察の両方を組み合わせているのでとても気に入っています。 - MoEレイヤーのボトルネック解析は、自然な最適化戦略をもたらします:メモリの読み書き回数をできるだけ減らすこと!前方方向方向の入力と後方方向方向の出力のグラデリングは、グループ化されたGEMMと同じくらい時間がかかることがあります。グループ化されたGEMM+オーバーラップメモリアクセスと融合、集め、計算することでレイヤー全体の速度を約2倍速くします。 - 専門家のルーティング用トップkの計算には驚くほど時間がかかることがあります。MoEレイヤー全体の約15〜20%です!標準的なtop-kインプルはradix top-kアルゴリズムを使っており、大きなkには優れていますが、小さなkには最適ではありません。Top-kはbitonic top-kアルゴリズムで書き直しましたが、pytorchのtop-kより20〜30倍速いこともあります! すべての主要なカーネルはCute-DSLで書かれているので、拡張やインストールが:Dやすいはずです。ホッパーのカーネルは出ていて、Blackwellのカーネルはほぼ完成しています。かつてはMoEモデルはトレーニングのハードウェア効率が2倍低かったので、Sonic-MOEがそれを変えてくれることを期待しています。