[Paper Review] 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, Phenomena53 references46 citations

TL;DR

FlowPM implements a GPU-accelerated, differentiable, distributed Particle-Mesh N-body solver in TensorFlow using a novel multi-grid, multiresolution pyramid scheme on Mesh-TensorFlow, achieving ~10x wallclock-speedups over Python FastPM and enabling differentiable forward modeling for cosmological inference.

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.

Motivation & Objective

• Enable fast, scalable cosmological N-body simulations on distributed GPUs using a differentiable PM framework.
• Develop and validate a novel multiresolution pyramid based multi-grid force estimation to reduce inter-process communication.
• Provide a Mesh-TensorFlow based architecture for distributed, model-parallel N-body computations.
• Demonstrate applicability to forward-modeling and cosmological field reconstruction with hybrid PM/ML forward models.

Proposed method

• Implement a FastPM-style PM evolution in TensorFlow for GPU acceleration and differentiability.
• Introduce a two-level multi-grid force estimation using multiresolution pyramids to split long-range and short-range forces across coarse global and fine local grids.
• Use a halo-exchange protocol with padding to enable slicewise operations on distributed grid chunks within Mesh-TensorFlow.
• Adopt Mesh-TensorFlow to map tensor computations onto a processor mesh with controlled distribution of grid dimensions.
• Construct and evaluate smoothed-downsample/upsample (REDUCE/EXPAND) operations via 3D convolutions with a bspline kernel for pyramid construction.
• Compare accuracy against the standard FastPM code and assess scaling on Cori GPUs with varying grid/mesh layouts.

Experimental results

Research questions

• RQ1 How can a PM cosmological N-body solver be implemented in a differentiable, GPU-accelerated framework?
• RQ2 Can a multi-grid, multiresolution pyramid scheme reduce distributed FFT communication while preserving accuracy?
• RQ3 What are the performance and scaling characteristics of FlowPM on large GPU meshes compared to a Python FastPM implementation?
• RQ4 Can FlowPM be effectively integrated into forward-modeling cosmological inference workflows (e.g., reconstruction with hybrid PM/ML components)?

Key findings

• FlowPM achieves sub-percent level accuracy relative to FastPM for the tested configurations (transfer function and cross-correlation within 0.01%).
• A two-level multi-grid pyramid scheme reduces distributed FFT communication by handling long-range forces on a coarse global grid and short-range forces on local grids.
• FlowPM scales on Cori GPUs with 8-V100 nodes, showing ~10x wallclock speed-up on average versus Python FastPM for comparable configurations.
• The Mesh-TensorFlow based implementation enables distributed model-parallel simulations with halo-exchange to manage boundary regions across process slices.
• The framework supports differentiable simulations suitable for forward modeling and reconstruction tasks, demonstrated via toy examples integrating neural network forward models.