表示
超・人工知能研究所 Super AI Research Lab トップへ超・人工知能研究所AI Research Lab / ぱそ太郎工事中

since
平成38年9月29日

← 検索結果に戻る

Frontierの通信経路に合わせ、言語モデルの重み・勾配・更新状態を分割する

Low-Bandwidth Partitioning

保存した記事を見るこのブラウザに保存します。無料会員で端末間の同期・フォルダ整理
原論文の図

Low-Bandwidth Partitioningの原論文から掲載した図です。細部はクリックして拡大できます。

原論文の図の説明を読む

Fig. 1: ZeRO-3 across two Frontier nodes.

出典:Low-Bandwidth Partitioning / arXiv ↗

概要

  • 重みを同一MI250X内の2GCD、勾配をノード内の8GCD、最適化状態を全GCDに分ける。
  • 著者らは200億パラメータのGPTモデルで、ZeRO++比1.71倍のGPU当たりTFLOPSと、384 GCDまでのスケーリング効率0.94を報告した。
  • 局所的な重み保持にはメモリ負担があるため、導入時は速度・最大メモリ使用量・学習損失を併せて測る。
補足:編集した模式図で仕組みを確認する

VISUAL EXPLAINER

図でつかむ、Low-Bandwidth Partitioningの仕組み

左から右へ読む
  1. 01学習用データ
  2. 02重みを2GCDで収集
  3. 03勾配をノード内で分配
  4. 04状態を用いて更新
  5. 05更新済み重み

原論文の説明をもとにした模式図です。処理の細部は省略しています。

研究の背景

大規模言語モデルの学習では、重みに加え、修正量を表す勾配や更新計算に使う状態が計算用メモリを占める。ZeRO-3はこれらを分割して保持するが、計算時には必要な値を集め直す通信が生じる。

Frontierでは、同じMI250X内の計算単位間、ノード内のGPU間、ノード間で通信帯域が異なる。そのため、何をどの範囲に分けるかが、通信時間と各計算単位のメモリ使用量に関わる。

手法

GCDはMI250X内の計算単位を指す。提案法はモデルの重みを同じMI250X内の2GCDに分け、前向き計算で必要な重みをそこで集める。後ろ向き計算用には二次的な重みの分割を保持し、数値を少ないビット数で表す「量子化」によって、その保存量を抑える。本文には二次分割を2GCDまたは8GCDに置く設計がある。

計算後の勾配は、重みをどの方向に直すかを示す値で、1ノード内の8GCDに分ける。更新に使う最適化状態は全GCDに分けて持ち、対応する勾配を同期して重みを更新する。例えばGPTモデルの学習では、各GCDが担当するデータを処理し、前向き計算、勾配計算、更新へ進む。

新規性

著者らの提案は、Frontierの通信帯域の違いに合わせ、重みを同一MI250X内、勾配をノード内、最適化状態を全GCDに分ける3段階の設計にある。重みの二次分割を量子化して保存する点も含む。

ZeRO++にも量子化と勾配通信の圧縮はある。確認できた本文で示される追加の工夫は、通信経路ごとに分割範囲を変え、局所的な重み保持によるメモリ負担を抑える組み合わせである。

従来手法との違い

ZeRO-3は重み・勾配・最適化状態を全計算単位に分けるため、重みの収集と勾配の分配がノードをまたぐ。ZeRO++は量子化やノード内の二次分割で通信量を減らすが、その分割にもメモリを使う。

提案法は重みを同一MI250X内の2GCD、勾配を1ノード内の8GCD、最適化状態を全GCDに分ける。二次的な重みの分割数には2と8の設計が記されており、後方計算時の重み収集範囲もそれに応じて変わる。

実験結果

著者らはFrontierで100億・200億パラメータのモデルを学習し、最大48ノード、192基のMI250X、384 GCDまで検証したと報告する。性能指標はGPU当たりのTFLOPS(1秒間に行う浮動小数点演算を兆単位で表す値)と、計算単位を増やした際のスケーリング効率である。

200億パラメータのGPTモデルでは、ZeRO++比でGPU当たりのTFLOPSが1.71倍、384 GCDまでのスケーリング効率が0.94だったという。ZeRO++自体についても、ZeRO-3比でGPU当たりのTFLOPSが最大40.5%増えたと著者らは報告する。確認できた本文の範囲では、学習データや各測定条件の詳細は確定できない。

応用の可能性

編集上の応用案 · 論文が実証した用途とは区別しています。

編集上の応用案:社内文書向け言語モデルの追加学習で、利用を認められた文書から学習用テキストを作り、更新済みモデルを得る工程に適用する。文章の続きを予測する課題なら、正解は元の文章から作る想定で、人手の正解ラベルやセンサー入力は使わない。

まず対象設備の計算単位間とノード間の通信経路を測り、重み・勾配・最適化状態の分割範囲を調整する。この用途で学習時間や業務上の回答品質が改善するかは、別途確かめる必要がある。

応用先の候補からNauhに相談できます。

導入時の検証

導入時の検証案 · 対象データでの再評価が必要です。

導入時の検証案:利用可能な学習用テキストを定め、同じモデル・データ量・学習設定でZeRO++の結果を基準として記録する。文章の続きを予測するなら、正解はテキストから作る。追加の人手ラベルやセンサーは使わず、対象設備の通信経路を調べて分割範囲を合わせる。

小規模な試行から計算単位を増やし、GPU当たりのTFLOPS、学習時間、最大メモリ使用量、学習損失を同条件で比べる。ノード間通信量も測定できる場合は記録し、速度とメモリ容量の両方を確認する。

限界と課題

著者らが示す設計上の制約は、通信を抑えるための局所的な重み保持がGCDのメモリを使う点である。ZeRO++についても、二次分割の追加により扱えるモデル規模が小さくなり得ると述べている。

編集上の検証事項:確認できた本文の範囲では、通信量分析の後半以降や、学習データと測定条件の全体を確認できない。導入先では最大メモリ使用量に加え、量子化を含む設定で学習結果が基準からどの程度変わるかを測る必要がある。

応用先の候補
研究内容と応用案に基づく分類です。業界での導入実績を示すものではありません。

出典

元のタイトル・要旨を確認する(英語)

Scaling Large Language Model Training on Frontier with Low-Bandwidth Partitioning

Scaling up Large Language Model(LLM) training involves fitting a tremendous amount of training parameters across a limited number of workers. However, methods like ZeRO-3 that drastically reduce GPU memory pressure often incur heavy communication to ensure global synchronization and consistency. Established efforts such as ZeRO++ use secondary partitions to avoid inter-node communications, given that intra-node GPU-GPU transfer generally has more bandwidth and lower latency than inter-node connections. However, as more capable infrastructure like Frontier, equipped with AMD GPUs, emerged with impressive computing capability, there is a need for investigations on the hardware topology and to develop targeted strategies to improve training efficiency. In this work, we propose a collection of communication and optimization strategies for ZeRO++ to reduce communication costs and improve memory utilization. In this paper, we propose a 3-level hierarchical partitioning specifically for the current 2nd ranked supercomputing cluster, Frontier, which aims at leveraging various bandwidths across layers of communications (GCD-GCD, GPU-GPU, and inter-node) to reduce communication overhead. For a 20B GPT model, we observe a 1.71x increase in TFLOPS per GPU when compared with ZeRO++ up to 384 GCDs and a scaling efficiency of 0.94 for up to 384 GCDs.

論文本体の取得範囲に基づく解説。AIが作成した未校閲の記事です。性能の数値は著者の評価条件に依存します。応用例と検証計画は編集上の提案です。解説更新:2026-10-01