大規模学習の通信を装置構成に合わせて効率化
HyperParallel-FSDP: Topology-Aware Fully Sharded Training with Layout-Driven Muon on Ascend SuperPods
この論文をやさしく読む
ひとことで言うと
多数の計算装置に分けたAI学習で、テンソル処理と通信の無駄を装置の接続構成に合わせて減らす仕組みです。
何に役立つ?
大規模モデルの分散学習で、通信待ちや分割管理の負荷を抑える実装の参考になります。
この研究の面白いところ
実運用の低負荷モードと厳密な検証モードに同じ分割計画を使い、速度と検証の両方を設計に組み込んでいます。
どこまで分かった?
数値はAtlas 900 A3 SuperPoDでの評価です。要旨のPearson相関は比較する量の詳細を明示しておらず、相関値だけから学習結果が完全一致するとまでは言えません。
v1のアブストラクトに基づくAI解説。日本語訳とは別に、用途の解釈を含みます。
アブストラクトの日本語訳
宣言的SPMDプログラミングは、テンソルの分割記述によって分散実行を制御し、並列化をモデルのコードから分離する。しかし評価対象のPyTorch DTensorスタックは、自動微分より下の層で各演算をディスパッチするため、ディスパッチとメタデータ処理の費用が繰り返し発生し、低コストな端から端までの検証経路も備えていない。既存のFSDPと分散Muonの実装も二階層のスーパーノード構成と整合していない。FSDPはパラメータの明示的なパックとアンパックに依存し、Muonの行列全体の直交化はパラメータ分割と衝突する。 本研究は、分散テンソルが分割の意味を表す必要があるのは自動微分より上のテンソルAPI境界だけであり、微分とカーネルは通常のテンソルを扱えることに着目する。この知見に基づきHyperParallel-FSDPを提案する。特徴は次の三つである。(1)二つのモードを持つDTensor実行。一つの分割計画を、レイアウトを一度だけ解決し定常時のディスパッチ負荷を持たない本番モードと、端から端までのメタデータ伝播、早期失敗検査、勾配等価性試験を行う検証モードの両方で利用する。(2)構成を考慮したFSDP。スーパーノード内のゼロコピー集団通信、スーパーノード間の融合リダクション、遅いリンクでの待機を避ける層横断の逆伝播パイプラインを用いる。(3)レイアウトに基づく分散Muon。分割から導く通信グループ、重複を除いた直交化、形状に基づき融合したNewton–Schulz反復を用いる。 Atlas 900 A3 SuperPoD上で、HyperParallel-FSDPは16ダイから384カード、768ランクまで拡張する。5050億パラメータの混合エキスパートモデルで毎秒42万1000トークンを維持し、FSDP通信はステップ時間の2.9%を占める。平均ステップ時間をPyTorch FSDP2に対して29.7%、Megatron DDPに対して25.5%削減し、1000ステップにわたるPearson相関は0.999997を上回る。分散Muonは競合システムに対して、プロファイラで測ったステップ時間を5.4〜16.0%改善する。ソースコードは https://atomgit.com/mindspore/hyper-parallel で公開されている。
v1の要旨から自動生成。本文の精読・人による確認は未実施。
- 初稿
- 2026-09-18(UTC)
- 最新改訂
- 2026-09-18 · v1
- 査読・掲載
- 査読状況未確認
更新履歴
- v1 2026-09-18 この版を読む
取得できた版を表示。版の更新は査読済みを意味しません。過去版の本文差分は未解析です。
原文の要旨
Declarative SPMD programming uses tensor sharding descriptions to drive distributed execution, separating parallelization from model code. However, the evaluated PyTorch DTensor stack dispatches every operator below autograd, incurring repeated dispatch and metadata costs, while lacking an inexpensive end-to-end validation path. Existing FSDP and distributed Muon implementations also mismatch two-tier supernode topologies: FSDP relies on explicit parameter packing and unpacking, and Muon's whole-matrix orthogonalization conflicts with parameter sharding. We observe that distributed tensors need only express sharding semantics at the tensor API boundary above autograd, allowing differentiation and kernels to operate on plain tensors. Based on this insight, we present HyperParallel-FSDP, featuring: (1) dual-mode DTensor execution, using one sharding plan for both a production mode with one-time layout resolution and no steady-state dispatch overhead, and a validation mode with end-to-end metadata propagation, fail-fast checks, and gradient-equivalence testing; (2) topology-aware FSDP, with zero-copy intra-supernode collectives, fused inter-supernode reduction, and a cross-layer backward pipeline that avoids waits on slow links; and (3) layout-driven distributed Muon, with sharding-derived communication groups, deduplicated orthogonalization, and shape-fused Newton-Schulz iterations. On Atlas 900 A3 SuperPoD, HyperParallel-FSDP scales from 16 dies to 384 cards (768 ranks), sustaining 421k tokens/s for a 505B-parameter MoE while FSDP communication uses 2.9% of step time. It reduces mean step time by 29.7% versus PyTorch FSDP2 and 25.5% versus Megatron DDP, with Pearson correlation above 0.999997 over 1,000 steps. Distributed Muon improves profiler step time by 5.4-16.0% over competing systems. Source code is available at https://atomgit.com/mindspore/hyper-parallel.
arXiv ID: 2609.21594 / 要約の誤りについて