# HexiScale: Facilitating Large Language Model Training over Heterogeneous Hardware
> [!info] Talk metadata
> - **会議:** [[MLSys2026]] Day 3 (May 20 / Wed)、Research Track Oral "LLM Training 3 & Model Compression"(17:00 PDT 開始、第3発表)
> - **登壇者:** Ran Yan(HKUST PhD student。研究関心は機械学習システムと最適化、特にヘテロGPU上での訓練・推論デプロイ、および sparse attention kernel 最適化)
> - **全著者:** Ran Yan¹*, Youhe Jiang¹*, Xiaonan Nie², Fangcheng Fu³, Bin Cui², Binhang Yuan¹(*: equal contribution、corresp. Binhang Yuan <
[email protected]>)
> - **所属:** ¹The Hong Kong University of Science and Technology (HKUST)、²Peking University、³Shanghai Jiao Tong University
> - **システム名:** HexiScale(ヘテロGPU上での大規模モデル訓練を非対称並列で実現するオープンソース訓練システム)
> - **関連リンク:** 連絡先
[email protected](スライド末尾)
> [!abstract] 概要(論文 PDF アブストラクト・忠実日本語訳)
> 大規模言語モデル(LLM)の訓練は計算集約的なタスクであり、通常は homogeneous な高性能 GPU を備えたデータセンタで実施される。本論文では代替アプローチとして、訓練計算を heterogeneous な GPU 上にデプロイし、ヘテロ資源の利用において高い柔軟性と効率を実現することを探究する。この目的のため、データ・パイプライン・テンソルモデル並列の枠組みにおける訓練計算の **非対称(asymmetric)分割**を柔軟にサポートできる新システム HexiScale を提案する。さらに、heterogeneous GPU 集合上への非対称分割された訓練計算の割当を**制約付き最適化問題**として定式化し、効率的な hierarchical graph partitioning アルゴリズムを提案する。本アプローチは利用可能な計算能力を最大限に活用しつつ、訓練計算を heterogeneous GPU 上へ効果的に割り当てる。HexiScale を state-of-the-art の homogeneous 系および heterogeneous 系の訓練システムと比較する。各種スケール(7B〜30B)の LLM を訓練したとき、実験結果は次を示す: (i) homogeneous GPU 上で動く SOTA の homogeneous ベースラインと比較し、**同一の理論 FLOPS** の heterogeneous GPU 上で HexiScale は **同程度**の性能を達成する。(ii) 同一の heterogeneous クラスタ上で動く SOTA の heterogeneous ベースラインと比較し、HexiScale は **1.5× 〜 2.4×** 高いスループットを実現する。
> [!note] 出典に関する注記
> - **発表者名の確定:** 文字起こしの自動生成では "Brandian"/"Ryan" と不鮮明だが、スライド末尾の連絡先 `
[email protected]` および筆頭著者 **Ran Yan**(HKUST)と整合するため、登壇者を Ran Yan と確定した(論文表紙・スライド表紙・スライド末尾が権威)。
> - **トークでの数値表現とアブストラクトの差異:** 発表者は hetero-aware ベースライン超えを「1.5×〜2.4×」ではなく **口頭で具体倍率を述べていない**箇所がある。論文アブストラクトの公式値は **1.5×〜2.4×**、スライド11の E2E 比較では **1.5×〜2.0×**(Llama 30B)を提示。本ノートはアブストラクト値を一次採用し、スライド値を併記する。
> - homo ベースライン比の平均スループットギャップ **0.83×**、最大 **1.01×** はスライド10が権威。case study の **1.6×**(symmetric 比、Llama-2 13B 仮想設定)は論文 §3.3 が権威。scheduler の **1.3×〜3.3×**(vs random graph partition)はスライド12が権威。
## 動機: ヘテロGPU訓練の意義
- **homogeneous な高性能 GPU は高価で希少(スライド2)。** LLM(Qwen、Llama 等)の訓練は数千 GPU の同時稼働を要するが、homogeneous な high-end GPU は小規模研究チームには取得困難。
- **既存資源が underutilize されている(スライド2)。** GPU 世代はおおよそ 2 年ごとに更新される(Turing 2018 → Ampere 2020 → Hopper 2022 → Blackwell 2024)が、旧世代 GPU は数年にわたり使われ続ける(例: K80 が AWS p2 インスタンスで今も提供、V100 も依然現役)。
- **Key question:** heterogeneous な GPU を用いて LLM 訓練で高性能を達成し、コスト効率の良い訓練と訓練の参入障壁の低減をコミュニティにもたらすには、どうすればよいか。散在する旧世代資源を統合し、訓練の barrier を下げることが狙い。
- HexiScale はオープンソース訓練システムであり、発表時点で **OSDI / SOSP 等のトップ会議に採択された複数論文の baseline として選定**されている(コミュニティが本研究を基盤・比較対象として採用している点を発表者が強調)。
## 2 つの課題(ヘテロ環境特有)
- **Challenge 1: heterogeneous GPU はハードウェア仕様が大きく異なる(スライド3、論文 §1)。**
- GPU 種別ごとに FP16 compute(peak TFLOPS)、HBM メモリ上限、通信能力が高い heterogeneity を示す(スライド3の散布図: RTX4090/L40/MXC500/910B/A100 ≈ 数百 TFLOPS、H100/H200 ≈ 1000 TFLOPS、MI3100X が最大)。
- 適切な管理がないと、能力の高い GPU が弱い GPU に合わせざるを得ず **underutilize される**。計算をハードウェアの能力に合わせて分割することが本質的に必要。
- **Challenge 2: GPU 間ネットワーク帯域が大きく異なる(スライド4、論文 §1)。**
- GPU-GPU ネットワーク接続は heterogeneous で、device topology 上のエッジは GPU ペア間の接続速度を表す。**NVLink 利用時は 300+ GB/s** に達するが、**PCIe / Ethernet のみだと 0.5 GB/s** まで低下し得る(スライド4の例: 300 GB/s / 16 GB/s / 0.5 GB/s が混在)。
- 速い接続が遅い接続に律速されて stall するのを防ぐため、異なる並列度へ割り当てる**ネットワーク資源の管理**が必要。
## 提案手法 1: 非対称並列サポート(Asymmetric Parallel Support)
HexiScale は **fully asymmetric parallelism**(データ・パイプライン・テンソルモデル並列の全枠組みで非対称分割)を導入する。symmetric setup(Megatron 等)と asymmetric setup(HexiScale)の違いは以下(スライド5、論文 §3.2、Figure 1)。
- **1 パイプライン内で、異なる PP stage が異なる TP degree を持てる**(symmetric では全 TP group が同一 degree 必須)。
- **異なるパイプライン間で、DP 同期を行う stage 同士が異なる TP degree を持てる**。例: 2 つの stage が data parallel 通信を互いに行いつつ、異なる TP degree を持つことが可能。これは Megatron 等の symmetric フレームワークでは非対応。
- **各 PP stage が異なる層数(number of layers)を持てる**(計算負荷の不均衡を補正)。
- **各 DP worker が異なる local batch size を持てる**(パイプライン間の実行時間バランスを取る)。
> [!note] 非対称分割の実装上の要点(論文 §3.2)
> - **計算負荷の非対称分割:** 各 data-parallel replica(= pipeline)に異なる batch size を割当て、各パイプライン並列通信グループを異なる TP degree・異なる割当層数で初期化。各 PP stage は近接 stage との通信レイテンシを最小化する leader GPU を選定し、forward では leader GPU が次 stage の leader へ activation を send、受信側が TP group 内へ broadcast して計算する(backward も同様に gradient を通信)。隣接 stage の TP degree が等しい場合は Megatron 流の scatter/gather に fall back。
> - **非対称 gradient 同期:** 異なる TP degree により同一層のパラメータ(と gradient)が異なるサイズに chunk される。vanilla data parallelism では同期できないため、最小 gradient size を特定し、より大きい gradient をその size に合わせて複数 chunk に分割し、DP worker の異なる subset 間で同じ通信量で同期する。
> - **システム実装:** ① gradient accumulation と activation recomputation で DP 通信コストとメモリフットプリントを償却、② FlashAttention-2(Hopper では FlashAttention-3 も)の API で TP + FlashAttention 付き transformer 層を生成、③ custom FSDP communication hook 登録で非対称 gradient 同期、④ NCCL ベースの communication group を data/tensor/pipeline 並列ごとに構成(NVIDIA GPU のみサポート)。
- **case study(論文 §3.1–3.3、Figure 1):** Llama-2 13B を A800-NVLink / 4090 / 3090 混在環境で訓練する仮想設定で、Megatron の対称戦略は推定 41.52s/iter に対し、HexiScale の非対称戦略は **25.55s/iter** に短縮し、**1.6× 高速**(= 1.6× 高いスループット)を達成。
## 提案手法 2: scheduling 問題の定式化と two-phase graph partitioning
### 問題定式化(論文 §4.1、スライド6)
- heterogeneous クラスタが与えられたとき、**訓練イテレーション時間を最小化する最適 placement** を求める。形式的には、最適並列実行プラン $\sigma^*$ を、メモリ制約下で通信コストと計算コストの和を最小化するものとして定める:
$\sigma^* = \arg\min_{\sigma}\ \mathrm{Comm\text{-}Cost}(\sigma) + \mathrm{Comp\text{-}Cost}(\sigma)\quad \text{s.t.}\quad \mathrm{Mem\text{-}Cumsum}(d) \le m_d\ \ \forall d \in \mathbf{D}$
ここで $d$ は GPU device、$m_d$ はそのメモリ上限、$\mathbf{D}$ は heterogeneous device 集合、$\sigma$ は並列構成。各 device $d$ のメモリ消費 $\mathrm{Mem\text{-}Cumsum}(d)$ がハードウェア上限 $m_d$ 以下でなければならない。
- 通信・計算コストは **cost model**(論文 Appendix B に詳細)で決定する。
- 探索空間は heterogeneous GPU により巨大であり、候補割当が指数的にスケールするため厳密最適化は **NP-hard**。Alpa / Galvatron 等の exhaustive search や dynamic programming ベースアルゴリズムは、ヘテロ環境では探索に長時間を要し(数百 GPU で数時間)、最適 placement を効率的に特定できない。HexiScale は **two-phase graph partitioning** ベースのヒューリスティックで効率的かつ効果的な近似解を求める。
### Phase-1: global graph partitioning(論文 §4.2、スライド7、Figure 2)
- device 集合をグラフ $G=(\mathbf{D},\mathbf{E})$ として抽象化。**頂点は計算能力(FLOPS)で重み付け**、**エッジは GPU-GPU 帯域**で重み付け。
- このグラフを $D_{dp}$ 個の **DP pool(各 pool が 1 つの DP worker = 1 パイプライン)** に分割する。$D_{dp}$-way multi-level graph partitioning(Hendrickson et al. 1995)に基づく **4 ステップ**:
1. **Coarsen:** Heavy Edge Matching (HEM) で高帯域接続の GPU を **supernode** にマージし、グラフを縮約(粗グラフは頂点数が半分程度)。クラスタが大規模でも効率的に処理できる。
2. **Partition:** recursive bisection で粗グラフを $D_{dp}$ グループに分割し、グループ間のネットワーク帯域を最小化(DP の境界で cut)。
3. **Project:** 分割を元グラフへ射影(project back)。
4. **Refine:** Kernighan-Lin アルゴリズムで境界を調整し partition 品質とバランスを向上(coarsening で失われた情報を回復)。
- **iterative optimization:** パイプライン数 $D_{dp}$ を enumerate して最適化。また、DP 通信用に帯域を最大化するか・PP 用に帯域を最小化するかを adaptive に切替える(パイプライン段数・batch size が大きいときは DP 帯域最小化+PP 帯域最大化、逆のときは DP 帯域最大化)。
### Phase-2: pipeline construction(論文 §4.3、スライド8、Figure 3)
各 DP worker(1 パイプライン)に割り当てられた GPU について、**3 ステップ**でパイプラインを構成する。
1. **Group GPUs for pipeline stages:** 割当 GPU を secondary graph に整理し、高帯域接続の GPU をさらにグループ分割。各グループが pipeline stage を構成し、グループ間は低帯域接続。
2. **Construct pipeline stages:** 各サブグループ(マシン)内で **intra-group strategy(TP/PP/層割当)** を cost model でシミュレーションし局所最適戦略を探索。
3. **Find pipeline stage order by greedy search:** 異なる GPU グループ間の遅い相互接続を避けるため、**top-$\tau$ greedy search** で pipeline stage の順序を permute し通信コストを最小化。
- **iterative optimization:** secondary graph partition のパラメータ $k_i$ を変えて異なる intra-group 戦略・stage 構成を生成し、性能を最大化。
- **batch / layer assignment:** batch は各パイプラインの推定訓練効率に比例配分、layer は推定計算効率に比例配分し、OOM 検出時は層数を減らしてメモリ余裕のある stage へ再割当て。各反復末に cost model でシミュレーションし、最小推定コストの戦略を選択(**simulation deviation < 2.9%**、論文 §5.4)。
## 評価(論文 §5)
### homogeneous ベースライン比較(スライド9/10、Figure 4/5)
- **ベースライン:** Megatron、Galvatron、FSDP。**比較スキーム:** 同一の総 FLOPS のもとで、heterogeneous GPU 上の HexiScale と homogeneous GPU 上のベースラインを比較(訓練は計算 bound のため総 FLOPS 固定が公平)。inter-machine 通信は両者とも Ethernet を使用。
- **実験設定(スライド9、3 つの hetero setting):**
- **Hetero-Setting 1:** 1×(8×3080Ti) + 1×(8×3090) + 3×(8×4090)。総 FLOPS +1.36%、メモリ 1.48× 小。Target homo: 2×(8×A800 PCIe-80G)。対象モデル Llama-2 7B / 13B。
- **Hetero-Setting 2:** 上記 + 1×(8×A800 NVLINK-80G)。総 FLOPS −1.59%、メモリ 1.14× 小。Target homo: 2×(8×A800 PCIe-80G)。Llama-2 7B / 13B。
- **Hetero-Setting 3:** 1×(8×3090) + 2×(4×3090) + 4×(8×4090) + 1×(8×A800 NVLINK-80G)。総 FLOPS −4.67%、メモリ 1.43× 小。Target homo: 4×(8×A800 PCIe-80G)。Llama 30B。
- **結果:** HexiScale は homogeneous ベースラインに匹敵するスループットを達成。homo の方が高いものの、**平均ギャップ 0.83×・最大 1.01×**(スライド10)。Llama-2 7B の例(スライド10): HexiScale 1.6 PFLOPS(Hetero-1)に対し Galvatron −46.8%、Megatron −95.2%。これは任意の GPU で訓練可能としつつ訓練障壁を下げられることを示す。FSDP は低速 inter-connect のクラスタには不適(homo Ethernet で −95% 超)。
### heterogeneity-aware ベースライン比較(スライド11、Figure 5)
- **ベースライン:** Metis、AMP、Espresso(および E2E 比較で Galvatron)。同一 heterogeneous クラスタ上で PFLOPS を比較。
- **結果:** HexiScale が一貫して上回る。より柔軟な非対称並列サポートと効果的な scheduler のため。スライド11(H800 + H20 混在)の AMP / Espresso 比: 16 H800 + 16 H20 で **1.79× / 1.78×**、16 H800 + 24 H20 で 1.73× / 1.72×、16 H800 + 32 H20 で 1.71× / 1.69×、Hetero-Setting 3 で 2.29× / 2.35×。E2E(Llama 30B)では Metis 比 **1.5×〜1.6×**、Galvatron 比 **1.7×〜2.0×**(latency breakdown でも HexiScale の comm/bubble/compute が最小)。
- 論文アブストラクトの公式総括は **1.5×〜2.4× 高スループット**(Metis は数百 GPU で探索が破綻、Megatron は対称制約で Hetero-Setting 3 の Llama 30B が OOM で動作不能)。
### scheduler の評価(スライド12、Figure 8)
- **effectiveness:** random graph partition と比較し、**早期収束**かつ平均 **1.3×〜3.3× 高い推定スループット**(Llama-2 7B Hetero-1 / Llama 30B Hetero-3)。
- **efficiency / scalability:** scheduling overhead は数千 GPU 規模でも manageable(coarsening 設計のおかげ)。アルゴリズム実行時間は **128 GPU で 6.7s、256 で 30.1s、512 で 96.0s、1024 で 372.1s**、推定 PFLOPS は **128 で 1.6、256 で 3.3、512 で 4.9、1024 で 6.9** と良好にスケール。
### ablation(論文 §5.3、Figure 6)
- **非対称並列の無効化:** 全 PP stage が同一 TP degree を強制され、最大 **23% 劣化**(Llama 30B)、平均 15% 劣化。
- **gradient accumulation の無効化:** 最大 **15% 劣化**(Llama 30B)、平均 12% 劣化(メモリ制約下の小 batch では DP 同期頻度が増えるため)。
## 結論(スライド13)
- **ヘテロGPU訓練の利点:** 異なる GPU 世代の散在資源を統合し、訓練の参入障壁をコミュニティに対して下げられる。
- **主要課題:** GPU のハードウェア仕様と低帯域ネットワーク接続の heterogeneity が、適切な管理なしには資源の underutilization を招く。
- **HexiScale の解:** (1) heterogeneous GPU 向けの柔軟な **非対称並列**の設計と実装、(2) 巨大な探索空間から効率的な構成を特定する **novel graph partition アルゴリズム**。
## Q&A
- **Q(Jensen, Georgia Tech):** 実際の graph partition を観察したとき、興味深いパーティションは見られたか。直感的には同じ種類の GPU が同じパーティションにグループ化されると期待するが、実際そうだったか、それとも別々にした方が良いケースもあったか。
- **A:** refinement は coarsening phase で失われた情報を回復するために使う。GPU を supernode にグルーピングして partition すると、元グラフに対しては partition が非効果的なことがある。refinement phase でそうした箇所をアルゴリズムで検出し、目的(GPU を boundary で分割=グループ間の tensor parallelism を最小化など)に partition を準拠させる。