[论文解读] FlowPM: Distributed TensorFlow Implementation of the FastPM Cosmological N-body Solver

Chirag Modi, François Lanusse+1 more|arXiv (Cornell University)|Oct 22, 2020

Galaxies: Formation, Evolution, Phenomena参考文献 53被引用 46

一句话总结

FlowPM 在 TensorFlow 中实现了一个 GPU 加速、可微分、分布式颗粒-网格 N体求解器，使用在 Mesh-TensorFlow 上的新型多网格、多分辨率金字塔方案，相比 Python FastPM 实现约 10x 的实际时钟时间加速，并实现用于宇宙学推断的可微分前向建模。

ABSTRACT

We present FlowPM, a Particle-Mesh (PM) cosmological N-body code implemented in Mesh-TensorFlow for GPU-accelerated, distributed, and differentiable simulations. We implement and validate the accuracy of a novel multi-grid scheme based on multiresolution pyramids to compute large scale forces efficiently on distributed platforms. We explore the scaling of the simulation on large-scale supercomputers and compare it with corresponding python based PM code, finding on an average 10x speed-up in terms of wallclock time. We also demonstrate how this novel tool can be used for efficiently solving large scale cosmological inference problems, in particular reconstruction of cosmological fields in a forward model Bayesian framework with hybrid PM and neural network forward model. We provide skeleton code for these examples and the entire code is publicly available at https://github.com/modichirag/flowpm.

研究动机与目标

• 在分布式 GPUs 上使用可微分 PM 框架实现快速、可扩展的宇宙学 N-body 模拟。
• 开发并验证基于新型多分辨率金字塔的多网格力估计，以减少进程间通信。
• 提供基于 Mesh-TensorFlow 的分布式、模型并行 N-body 计算架构。
• 展示在前向建模和带混合 PM/ML 的宇宙学场重建中的适用性。

提出的方法

• 在 TensorFlow 中实现 FastPM 风格的 PM 演化以实现 GPU 加速和可微分性。
• 引入使用多分辨率金字塔的两级多网格力估计，将远程力分配到粗 global 网格，短程力分配到局部细网格。
• 使用带填充的 halo 交换协议，以使 Mesh-TensorFlow 内分布网格块上的切片操作成为可能。
• 采用 Mesh-TensorFlow 将张量计算映射到处理器网格，并对网格维度的分布进行控制。
• 通过 3D 卷积使用 bspline 内核构建并评估平滑下采样/上采样（REDUCE/EXPAND）操作用于金字塔构建。
• 对 accuracy 与标准 FastPM 代码进行比较，并在 Cori GPU 上使用不同网格/网格布局评估扩展性。

实验结果

研究问题

• RQ1 如何在可微分、GPU 加速框架中实现一个 PM 宇宙学 N-body 求解器？
• RQ2 多网格、跨尺度金字塔方案在保持准确性的同时是否能减少分布式 FFT 的通信？
• RQ3 与 Python FastPM 实现相比，FlowPM 在具有大规模 GPU 网格时的性能和扩展性特征是什么？
• RQ4 FlowPM 能否有效集成到前向建模的宇宙学推断工作流中（例如与混合 PM/ML 成分的重建）？

主要发现

• FlowPM 相对于 FastPM 在测试配置上达到亚百分比级别的准确性（传递函数和互相关性在 0.01% 内）。
• 两级多网格金字塔方案通过在粗全局网格处理长程力、在局部网格处理短程力，降低分布式 FFT 通信。
• FlowPM 在 Cori GPU 的 8-V100 节点上扩展性良好，对比相应配置时平均实现约 10x 的 wallclock 加速，相较于 Python FastPM。
• 基于 Mesh-TensorFlow 的实现实现了分布式模型并行仿真，并通过 halo 交换来管理跨进程切片的边界区域。
• 该框架支持适用于前向建模与重建任务的可微分仿真，通过将神经网络前向模型整合在 toy 示例中得以演示。