250918-02
出典: huawei.com

Huaweiは、最大100万個のAIチップを1台のコンピューターとして動かすことを目指す計算アーキテクチャ「Peerium」を披露しました。同社は最大40兆パラメータのAIモデルの学習を想定しており、演算を分担するチップ間の通信を独自技術で支える構成です。計算能力を追加するハードウェアモジュールである計算カードを25万6000枚接続した「SuperCluster」は、配備と試験が進んでいます。

チップの階層的な分担とメモリアクセスを組み合わせる

Peeriumを構成するのは、入れ子状の並列処理、統一メモリアドレス方式、機器同士を直接つなぐピアツーピア接続の3要素です。入れ子状の並列処理では、処理を小さく分け、階層化したチップのグループへ割り当てて同時に実行します。

統一メモリアドレス方式では、プロセッサーが自身のサーバーに接続されたメモリだけでなく、別のサーバーのメモリにもアクセスできます。ピアツーピア接続は、機器同士が仲介を減らして直接通信する仕組みです。Huaweiは、これらを組み合わせることで、100万個規模までプロセッサーを増やして協調動作させられるとしています。

接続を担う独自技術「UnifiedBus」は、CPU、AI処理用のNPU、メモリ、SSD、インターフェースカード、スイッチを結ぶ高速バスです。計算、記憶、ネットワークの各機器を共通のプロトコル、つまり通信の取り決めで接続します。

Huaweiの輪番会長エリック・シュー氏は、UnifiedBusによるラック間の通信帯域は最大800GB/秒に達するとしています。同氏は、NVIDIAのNVL72ではラック内にNVLink、ラック間にInfiniBandという別々の接続技術を使い、ラック間の帯域は200GB/秒になると比較しました。

40兆パラメータの学習を想定、供給は段階的に拡大

Peeriumを採用する計算システム「Atlas 950 SuperPoD」は、HuaweiのAIアクセラレーターチップ「Ascend 950DT」を8000個以上搭載します。SuperClusterは、複数のSuperPoDを接続した、より大規模なAI計算システムです。

Huaweiの半導体設計部門HiSiliconの首席科学者リャオ・ヘン氏は、今後2年間に中国の先端AI研究所6～7組織が、10兆～40兆パラメータの基盤言語モデルの学習を目指すと予測しました。同氏は、こうしたモデルの規模に対応するメモリ容量と計算基盤が必要になるとの考えを示しています。

Ascend 950DTを使うSuperPoDは主にAI学習向けで、試験が進められています。シュー氏は、試験結果を良好と評価し、大規模供給を年末から翌年初めに開始する予定を示しました。また、中国国内の需要を満たす供給能力が足りないため、海外市場への本格展開は計画せず、一部の国での試験と少量の供給にとどめています。

学習方法によって変わるチップ間の通信

Peeriumが扱う大規模AI学習では、複数の計算機で処理を分担する方法に応じて、やり取りするデータや待ち時間が変わります。データ並列では、各計算プロセスがモデル全体を持ち、異なる学習データを処理します。学習の各ステップで、重みを更新する方向を表す「勾配」をAllReduceという集団通信で集約し、各プロセスで共有する仕組みです。

モデル自体を分ける方法には、層と層の間で分担するパイプライン並列と、1つの層の内部を分担するテンソル並列があります。パイプライン並列では、処理の受け渡しに伴う待ち時間である「パイプラインバブル」が効率を下げます。テンソル並列では層ごとに通信が必要になり、演算と通信を重ねて実行できないことがボトルネックです。

スーパーコンピューター「富岳」の分散学習でも、こうした通信と待ち時間が性能に影響しています。東京工業大学の横田理央氏らの研究では、演算と通信の高速化前はテンソル並列の通信時間が目立ち、高速化後はパイプラインバブルによる待ち時間が大きく現れました。

40兆パラメータをBF16で保存すると重みだけで80TB

Huaweiが想定する40兆パラメータという規模は、モデルの重みを保存するメモリにも関わります。重みを16ビット形式のBF16で保存する場合、1パラメータは2バイトなので、40兆個の重みの保存容量は10進表記で80TBです。学習時には、重みに加えて勾配、重みの更新を管理するオプティマイザーの状態、計算途中の値などもメモリを使います。

パラメータの使い方はモデルの構造でも異なります。全結合型のDenseモデルは各トークンの処理で全パラメータを使う一方、複数の処理ネットワークを持つMoEでは、その一部を選択して動かします。MoEでは全体の重みを保持するメモリ容量が総パラメータ数に関わり、各トークンを処理する演算量は、その際に動くパラメータ数によって変わる仕組みです。

既存GPUでも別の機器のメモリへ直接アクセス

Peeriumが取り入れる別のサーバーのメモリへのアクセスに関連する仕組みは、既存のGPU向け技術にもあります。NVIDIAの「GPUDirect RDMA」は、PCI Expressを使い、GPUとネットワーク機器などの間で直接データを交換する経路を設けます。機器間でメモリを読み書きするために利用するのは、ドライバーによるアドレス変換や対応付けです。

GPUから通信を開始する「NVSHMEM」の研究では、PCIeやNVLinkで接続したGPUに加え、InfiniBandを介した通信をGPU内の実行プログラムから発行する仕組みが実装されました。共有対象として確保したメモリ領域を、別の処理プロセスから読み書きする方式です。