[論文レビュー] Large-Scale Discrete Fourier Transform on TPUs

Tianjian Lu, Yifan Chen+3 more|arXiv (Cornell University)|Feb 9, 2020

Parallel Computing and Optimization Techniques参考文献 34被引用数 11

ひとこと要約

本論文は、TPUクラスタ上で大規模3次元離散フーリエ変換（DFT）を実行するための2つの並列アルゴリズム—Kronecker積と密行列乗算に基づくKDFTと、位相補正を施したCooley-Tukeyに基づくFFT—を提案する。TensorFlowで実装され、TPUの行列乗算の強みと1回のシャッフル通信スキームを活用しており、2048コアのフルTPU Podで8192³ DFTを8.3秒で実行し、高い並列効率と低い通信オーバーヘッドを達成した。

ABSTRACT

In this work, we present two parallel algorithms for the large-scale discrete Fourier transform (DFT) on Tensor Processing Unit (TPU) clusters. The two parallel algorithms are associated with two formulations of DFT: one is based on the Kronecker product, to be specific, dense matrix multiplications between the input data and the Vandermonde matrix, denoted as KDFT in this work; the other is based on the famous Cooley-Tukey algorithm and phase adjustment, denoted as FFT in this work. Both KDFT and FFT formulations take full advantage of TPU's strength in matrix multiplications. The KDFT formulation allows direct use of nonuniform inputs without additional step. In the two parallel algorithms, the same strategy of data decomposition is applied to the input data. Through the data decomposition, the dense matrix multiplications in KDFT and FFT are kept local within TPU cores, which can be performed completely in parallel. The communication among TPU cores is achieved through the one-shuffle scheme in both parallel algorithms, with which sending and receiving data takes place simultaneously between two neighboring cores and along the same direction on the interconnect network. The one-shuffle scheme is designed for the interconnect topology of TPU clusters, minimizing the time required by the communication among TPU cores. Both KDFT and FFT are implemented in TensorFlow. The three-dimensional complex DFT is performed on an example of dimension $8192 imes 8192 imes 8192$ with a full TPU Pod: the run time of KDFT is 12.66 seconds and that of FFT is 8.3 seconds. Scaling analysis is provided to demonstrate the high parallel efficiency of the two DFT implementations on TPUs.

研究の動機と目的

• 科学計算ワークロード向けに、TPUクラスタ上で高性能で大規模な離散フーリエ変換（DFT）を可能にすること。
• Kronecker積（KDFT）とCooley-Tukey FFT（FFT）を用いてDFTを密行列演算として定式化することで、TPUの行列乗算の強みを活用すること。
• TPUのインターコネクトトポロジーに適合した通信効率の高いデータ分割とコア間通信スキームを設計すること。
• 数千のTPUコアにわたる3次元DFTにおいて、強いスケーリングと高い並列効率を達成すること。
• 実世界の科学的応用に向けたTPU上でのDFT高速化の実現可能性とパフォーマンスを実証すること。

提案手法

• Kronecker積を用いて3次元DFTを入力データとバーデモンド行列との行列乗算として定式化（KDFT）し、非一様入力の直接処理を可能にする。
• 位相補正を施したCooley-TukeyアルゴリズムによるFFTの実装により、計算量をO(N²)からO(N log N)に低減する。
• データ分割を適用して、各TPUコア内に密行列乗算を局所化し、クラスタ全体で完全な並列処理を実現する。
• 同じインターコネクト方向に沿って隣接するTPUコア間で同時に送信と受信が可能な1回のシャッフル通信スキームを採用し、通信遅延を最小限に抑える。
• einsumやcollective_permuteなどのハードウェア最適化演算を活用して、TPU Pod上でポータブルかつスケーラブルに実行可能なTensorFlowへの両アルゴリズムのマッピング。
• 128 GiBの高帯域幅メモリと直接チップ間接続を活用し、大規模なデータ移動におけるボトルネックを最小限に抑える。

実験結果

リサーチクエスチョン

• RQ1 DFTのKronecker積による定式化は、行列乗算プリミティブを用いてTPUクラスタ上で効率的に並列化可能か？
• RQ2 大規模3次元問題において、FFTベースのDFTは、直接行列乗算（KDFT）による定式化と比較して、TPU上でどの程度性能に優れるか？
• RQ3 1回のシャッフル通信スキームは、大規模DFT計算におけるTPU間通信オーバーヘッドをどの程度低減できるか？
• RQ4 2048コアにまで達するフルTPU Pod上での3次元DFTの強スケーリング特性はいかがなものか？
• RQ5 TPUベースのDFT実装は、科学計算ワークロードにおいて高い並列効率と低遅延を達成できるか？

主な発見

• 2048コアのフルTPU Pod上での3次元FFTは、8192×8192×8192 DFTを8.30秒で実行し、同じ問題に対してKDFTが12.66秒を要したのと比較して顕著に優れた性能を示した。
• 2048×2048×2048問題では、KDFTとFFTの両者ともほぼ同一の計算時間（0.120秒と0.118秒）を記録しており、スケールが大きくなるとTPUの行列乗算性能がアルゴリズム的複雑度の差を上回ることが示された。
• 強いスケーリング解析により、2048³問題における3次元FFTは128コアまでほぼ線形の性能スケーリングを示し、理想的に近いスループット向上が観察された。
• 128コアを超えてスケーリングする際、all_to_allおよびcollective_permute演算に起因する通信時間がボトルネックとなった。これは、大規模な通信パターンにおけるさらなる最適化の必要性を示唆している。
• 3次元KDFTおよびFFT実装は高い並列効率を示し、アルゴリズム的複雑度で正規化した場合、計算時間は問題サイズにほぼ線形に比例した。
• 1回のシャッフル通信スキームは、同じインターコネクト方向に沿って同時に送信と受信を可能にすることで、TPU間通信遅延を効果的に低減した。