[논문 리뷰] Large-Scale Discrete Fourier Transform on TPUs

Tianjian Lu, Yifan Chen외 3명|arXiv (Cornell University)|2020. 02. 09.

Parallel Computing and Optimization Techniques참고 문헌 34인용 수 11

한 줄 요약

이 논문은 TPU 클러스터에서 대규모 3D 이산 푸리에 변환을 위한 두 가지 병렬 알고리즘—Kronecker 곱과 희소 행렬 곱셈 기반의 KDFT와 Cooley-Tukey 알고리즘을 기반으로 한 FFT(단계 조정 포함)—을 제시한다. Tensorflow로 구현된 이 알고리즘들은 TPU의 행렬 곱셈 성능과 한 번의 셔플 통신 방식을 활용하여, 전체 2048코어 TPU 팟에서 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 곱법을 사용하여 DFT를 밀도 있는 행렬 연산으로 재구성함으로써 TPU의 행렬 곱셈 성능을 극대화하고 Cooley-Tukey FFT(FFT)를 적용하기 위해.
• TPU의 인터커넥트 토폴로지에 맞는 통신 효율적인 데이터 분할 및 코어 간 통신 방식을 설계하기 위해.
• 수천 개의 TPU 코어에 걸쳐 3D DFT의 강한 스케일링과 높은 병렬 효율성을 달성하기 위해.
• 실제 과학 응용 분야에서 TPU를 활용한 DFT 가속화의 가능성과 성능을 입증하기 위해.

제안 방법

• Kronecker 곱을 활용하여 입력 데이터와 바르모인드 행렬 간의 행렬 곱셈으로 3D DFT를 재구성함으로써 비균일 입력 데이터를 직접 처리할 수 있도록 한다(KDFT).
• 단계 조정 기반의 Cooley-Tukey 알고리즘을 구현하여 계산 복잡도를 O(N²)에서 O(N log N)으로 감소시킨다.
• 밀도 있는 행렬 곱셈을 개별 TPU 코어 내부에 국한시키기 위해 데이터 분할을 적용함으로써 클러스터 전반에 걸친 완전한 병렬성을 달성한다.
• 이중 인터커넥트 방향에서 이웃한 TPU 코어 간에 동시에 보내기 및 받기 작업을 수행할 수 있도록 한 번의 셔플 통신 방식을 사용하여 통신 지연을 최소화한다.
• einsum 및 collective_permute와 같은 하드웨어 최적화 연산을 활용하여, TPU 팟에서 이식성 있고 확장 가능한 실행을 위한 TensorFlow에 KDFT와 FFT를 매핑한다.
• 고대역폭 메모리(128 GiB)와 직접적인 칩 간 인터커넥트를 활용하여 대규모 데이터 이동에서의 병목 현상을 최소화한다.

실험 결과

연구 질문

• RQ1 DFT의 Kronecker 곱 표현 방식이 TPU 클러스터에서 행렬 곱셈 프리미티브를 활용하여 효율적으로 병렬화될 수 있는가?
• RQ2 대규모 3D 문제에서 TPU에서 FFT 기반 DFT의 성능이 직접 행렬 곱셈(KDFT) 방식보다 어떻게 다를 것인가?
• RQ3 한 번의 셔플 통신 방식이 대규모 DFT 계산에서 TPU 간 통신 오버헤드를 얼마나 줄이는가?
• RQ4 2048개 코어까지 확장 가능한 전체 TPU 팟에서 3D DFT의 강한 스케일링 특성은 어떠한가?
• RQ5 TPU 기반 DFT 구현은 과학 계산 워크로드에서 높은 병렬 효율성과 낮은 지연을 달성할 수 있는가?

주요 결과

• 전체 2048코어 TPU 팟에서 8192×8192×8192 DFT에 대해 3D FFT는 8.30초의 실행 시간을 기록하여 동일한 문제에서 KDFT의 12.66초보다 뚜렷이 뛰어난 성능을 보였다.
• 2048×2048×2048 문제에서는 KDFT와 FFT가 거의 동일한 계산 시간(0.120s 및 0.118s)을 기록하여, 대규모에서 TPU의 행렬 곱셈 성능가 알고리즘 복잡도의 차이를 초월함을 시사했다.
• 강한 스케일링 분석 결과, 2048³ 문제에서 3D FFT는 128코어까지 거의 선형적인 성능 향상을 보였으며, 이상적인 스피드업이 관찰되었다.
• 128코어를 초과하여 확장할 경우, all_to_all 및 collective_permute 연산에 의해 주로 결정되는 통신 시간이 병목이 되었으며, 이는 대규모 통신 패tern에 대한 추가 최적화가 필요함을 시사했다.
• 3D KDFT와 FFT의 구현은 높은 병렬 효율성을 보였으며, 알고리즘 복잡도에 따라 정규화된 경우 계산 시간이 문제 크기와 약간의 선형 비례를 보였다.
• 한 번의 셔플 통신 방식은 동일한 인터커넥트 방향을 따라 동시에 보내기 및 받기 작업을 가능하게 하여 TPU 간 통신 지연을 효과적으로 최소화했다.