PM++: Scaling Differentiable Cosmological Simulators with Dynamic Particle Ownership

Large differentiable particle-mesh simulations support cosmological inference and simulator-coupled learning, but their evolving states can exceed single-accelerator memory. Distribution with fixed particle ownership remains memory-limited because deposition halos must cover the maximum cumulative displacement. We present $\texttt{PM++}$, a JAX-native simulator that distributes particles and meshes across GPUs, migrates authoritative particle records between spatial slabs, and preserves static array shapes. A reversible discrete adjoint differentiates through routing, cloud-in-cell operations, halo exchange, and distributed Fourier transforms without storing the full trajectory. Relative to two H200 GPUs, eight GPUs give $\texttt{PM++}$ $2.73\times$ forward and $2.66\times$ differentiated speedups at $1024^3$ particles, with near-inverse per-GPU memory scaling. On eight GPUs, $\texttt{PM++}$ completes forward and differentiated runs at $1800^3$ and $1536^3$ particles, respectively. In a $100\ \mathrm{Mpc}\ h^{-1}$ box, a fixed 128-cell halo yields a $3.4634\%$ mass deficit for displacements of approximately 410 cells, whereas dynamic ownership conserves mass. $\texttt{PM++}$ thus expands the volume--resolution regime available to field-level inference and differentiable correction models.