[论文解读] High Performance Monte Carlo Simulation of Ising Model on TPU Clusters

Kun Yang, Yifan Chen+3 more|arXiv (Cornell University)|Mar 27, 2019

Theoretical and Computational Physics参考文献 17被引用 8

一句话总结

本文提出了一种使用 TensorFlow 在 Cloud TPU 集群上对二维伊辛模型进行高性能蒙特卡洛模拟的方法，利用 TPU 的矩阵运算能力和高速互连网络，实现了比以往基准测试高 250% 的多核性能。该实现仅用一个 Jupyter Notebook 完成，采用 bfloat16 精度保持了准确性，并在最多 2048 个核心下展现出良好的线性加速性能。

ABSTRACT

Large-scale deep learning benefits from an emerging class of AI accelerators. Some of these accelerators' designs are general enough for compute-intensive applications beyond AI and Cloud TPU is one such example. In this paper, we demonstrate a novel approach using TensorFlow on Cloud TPU to simulate the two-dimensional Ising Model. TensorFlow and Cloud TPU framework enable the simple and readable code to express the complicated distributed algorithm without compromising the performance. Our code implementation fits into a small Jupyter Notebook and fully utilizes Cloud TPU's efficient matrix operation and dedicated high speed inter-chip connection. The performance is highly competitive: it outperforms the best published benchmarks to our knowledge by 60% in single-core and 250% in multi-core with good linear scaling. When compared to Tesla V100 GPU, the single-core performance maintains a ~10% gain. We also demonstrate that using low precision arithmetic---bfloat16---does not compromise the correctness of the simulation results.

研究动机与目标

• 探索使用 AI 加速器（如 Cloud TPU）进行科学计算的可行性，超越深度学习应用。
• 在 TPU 上使用 TensorFlow 实现高性能、可扩展且可读性强的二维伊辛模型蒙特卡洛模拟。
• 评估低精度算术（bfloat16）对统计物理中蒙特卡洛模拟精度的影响。
• 展示模拟在数千个 TPU 核心上的强线性扩展性能。

提出的方法

• 在 Cloud TPU 上使用 TensorFlow 实现基于 SIMD 的分布式马尔可夫链蒙特卡洛（MCMC）算法，用于二维伊辛模型。
• 利用 TPU 的高带宽内存和二维环形网格互连结构，实现高效的数据传输和低延迟通信。
• 使用 tf.nn.conv2d 高效计算最近邻能量贡献，替代批量矩阵乘法，以更好地利用 MXU 的计算能力。
• 通过松散打包、密集打包和超密集打包配置，优化工作负载在 TPU 核心上的分布，以最大化资源利用率。
• 使用 TensorFlow r1.15 和 TPU v3 硬件，在代码复杂度极低的前提下实现高性能。
• 通过 Google Colab 在单个 Jupyter Notebook 中运行完整模拟，仅需极少设置，即可在 TPU 上实现交互式科学计算。

实验结果

研究问题

• RQ1 Cloud TPU 集群能否被有效用于深度学习之外的高性能科学模拟？
• RQ2 基于 TensorFlow 的 TPU 上伊辛模型模拟性能与已发表的 GPU 和 CPU 基准测试相比如何？
• RQ3 使用 bfloat16 精度是否会损害统计物理中蒙特卡洛模拟的正确性？
• RQ4 该模拟在数千个 TPU 核心上的扩展性能在多大程度上是线性的？
• RQ5 复杂的分布式科学算法能否在 TensorFlow 中清晰且高效地表达，同时不牺牲性能？

主要发现

• 基于 TPU 的实现相比最佳已发表基准，在单核配置下提升 60%，在多核配置下提升 250%，并展现出极佳的线性扩展性能。
• 在拥有 32,768 个核心的完整 TPU v3 Pod 上，模拟达到每纳秒 20,460.92 次自旋翻转，表明在 1,000 个核心以内具有强大的扩展能力。
• 在超过 1,000 个核心后，通信开销变得显著，限制了强扩展测试中的进一步扩展效率。
• 使用 bfloat16 精度不会损害模拟结果的正确性，保持了物理精度。
• 整个模拟仅用一个清晰可读的 Jupyter Notebook 实现，可通过 Google Colab 以极少设置完成执行。
• 通过使用 tf.nn.conv2d 替代批量乘法的优化版本，性能相比初始实现提升了约 80%。