$\texttt{HEXA-MoE}$：計算の冗長性をゼロに抑えた、効率的でヘテロジニアス対応のMoE高速化手法

Core Concepts

本稿では、計算の冗長性をゼロに抑え、データ中心とモデル中心の両方の構成に対応し、ヘテロジニアス環境での利用も可能な、効率的なMoE高速化フレームワークであるHEXA-MoEを提案する。

Abstract

HEXA-MoE: 計算の冗長性をゼロに抑えた、効率的でヘテロジニアス対応のMoE高速化手法

本稿では、Transformerモデルのパラメータを効率的にスケールアップするための実用的手法であるMixture-of-Experts (MoE) の高速化フレームワーク、HEXA-MoEについて解説する。

Transformerは、自然言語処理、コンピュータビジョン、マルチモーダル学習など、幅広い機械学習タスクにおいて標準的なアーキテクチャとなっている。モデルのパラメータ数を増やすことで学習能力と汎化能力が向上することが知られているが、計算量の増加が課題となる。MoEは、計算のオーバーヘッドをサブ線形に抑えながらパラメータ数をスケールアップするための効果的なアプローチとして注目されている。
従来のMoEモデルは、分散デバイス上でのエキスパート並列処理で構築されることが一般的であった。しかし、この手法は、均質なデバイスへの展開を前提としており、通信オーバーヘッドと計算の冗長性が課題となっていた。

Key Insights Distilled From

$\texttt{HEXA-MoE}$: Efficient and Heterogeneous-aware MoE Acceleration with ZERO Computation Redundancy

by Shuqing Luo,... at arxiv.org 11-05-2024

https://arxiv.org/pdf/2411.01288.pdf

Deeper Inquiries

MoEの高速化は、Transformer以外のモデルアーキテクチャにも適用できるか？

MoEの高速化は、Transformer以外にも、計算コストの高い層を疎に活性化するという特性を持つモデルアーキテクチャに適用できる可能性があります。
例えば、以下のケースが考えられます。

巨大な出力空間を持つモデル: 例えば、言語モデルにおいて、語彙数が非常に大きい場合、出力層の計算コストが非常に高くなります。この場合、語彙を複数のエキスパートに分割し、入力に応じて少数のエキスパートのみを活性化するMoEを適用することで、計算コストを抑えながら、巨大な出力空間を扱うことができます。
多様なタスクを学習するモデル: 例えば、画像認識と自然言語処理を同時に行うマルチモーダルモデルの場合、タスクごとに異なるエキスパートを用意し、入力に応じて適切なエキスパートを選択することで、各タスクに特化した学習が可能になります。
ただし、MoEの適用には、エキスパートのルーティングやロードバランスなど、解決すべき課題も存在します。Transformer以外のモデルに適用する場合は、これらの課題に対する適切な対策を検討する必要があります。

HEXA-MoEは、計算の冗長性を完全に排除しているが、計算の精度に影響はないのか？

HEXA-MoEは、計算の冗長性を排除するために、Expert-Specific Operatorsを用いて、従来のGeMM (General Matrix Multiplication) ベースの計算方式から、In-placeな計算方式に変更しています。
この変更により、計算精度の面では、以下の点が懸念されます。

AtomicAddの利用: Top-kルーティングにおけるメモリ効率を向上させるため、HEXA-MoEはAtomicAdd操作を利用しています。AtomicAddは並列計算における競合状態を引き起こす可能性があり、計算精度に微小な影響を与える可能性があります。
数値の丸め誤差: In-placeな計算方式では、計算の順序や中間結果の保持方法が変わるため、従来の計算方式と比較して、数値の丸め誤差がわずかに異なる可能性があります。
しかし、論文では、HEXA-MoEが従来のMoEライブラリと比較して、同等の精度を達成することを実験的に確認しています。これは、HEXA-MoEが計算精度の影響を最小限に抑えるように設計されているためと考えられます。
ただし、モデルやタスクによっては、計算精度の影響が大きくなる可能性も否定できません。そのため、HEXA-MoEを適用する際には、精度検証を行うことが重要です。

ヘテロジニアス環境におけるMoEの高速化は、今後ますます重要になるか？

ヘテロジニアス環境におけるMoEの高速化は、以下の理由から、今後ますます重要になると考えられます。

大規模化の進展: 深層学習モデルは、今後も大規模化が進むと予想されます。大規模なMoEモデルを効率的に学習するためには、GPUやTPUなど、異なる種類のデバイスを組み合わせたヘテロジニアス環境での高速化が不可欠になります。
ハードウェアの多様化: GPUやTPUなどのアクセラレータは、性能や価格帯が異なる様々な種類が開発されています。これらの多様なハードウェアを効率的に活用するためには、ヘテロジニアス環境に対応した高速化技術が求められます。
コスト削減: 最新の高性能GPUは高価であるため、旧世代のGPUや性能の異なるGPUを組み合わせたヘテロジニアス環境は、コスト削減の観点からも重要です。
HEXA-MoEは、データ並列とテンソル並列を組み合わせることで、ヘテロジニアス環境においても効率的なMoE計算を実現しています。これは、今後のMoEの高速化において、重要な方向性の一つを示していると言えるでしょう。
さらに、デバイスの計算能力に応じた最適なワークロード分配や、通信コストを考慮したエキスパート配置など、ヘテロジニアス環境におけるMoE高速化には、更なる研究開発が必要です。