One-sentence Summary

Researchers from the University of Nebraska–Lincoln, the University of Rochester Medical Center, and the Holland Computing Center present a scalable multi-GPU framework for three-dimensional multicellular growth simulation that combines GPU acceleration, spatial binning, domain decomposition, and workload-aware partitioning with an RNN-based load-balancing controller, which observes recent per-rank execution times and partition states, learns residual corrections to a reactive boundary-adjustment rule, and is trained offline in a differentiable surrogate with randomized workload dynamics and no execution traces; the controller reduces mean global imbalance from $11.3\%$11.3% under static partitioning to $3.5\%$3.5%, lowers end-to-end runtime by $9.0\%$9.0% relative to static partitioning, and reduces slice migration by $7.7\times$7.7× compared with the reactive baseline.

Key Contributions

• A scalable multi-GPU framework for three-dimensional multicellular growth simulation combines GPU acceleration, spatial binning, domain decomposition, and workload-aware partitioning; GPU acceleration with spatial binning reduces element-interaction time by roughly three orders of magnitude over a serial CPU baseline, with binning providing a further 9.2× speedup over a GPU all-pairs implementation at 6.3 million elements.
• An RNN-guided load-balancing controller learns residual corrections to a reactive boundary-adjustment rule from recent per-rank execution times and partition states, and it is trained offline in a differentiable surrogate of the load-balancing loop without measured execution traces. In end-to-end 435K-cell simulations on 8 GPUs, it reduces mean global imbalance from 11.3% under static partitioning to 3.5%, compared with 3.8% for the reactive rule, lowers total runtime by 9.0% relative to static partitioning, and migrates 7.7× fewer slices than the reactive baseline with comparable runtime.
• A representative embryonic epidermal development use case involving proliferation, differentiation, and tissue reorganization demonstrates support for spatially and temporally evolving workloads and a substantially improved balance-migration tradeoff.

Introduction

Multicellular tissue development emerges from local cell behaviors such as growth, division, migration, and adhesion, and three-dimensional cell-based models make these mechanisms explicit and controllable. Subcellular element models provide finer mechanical detail by representing each cell with multiple interacting elements, but this increases force-evaluation cost significantly, and the workload becomes spatially nonuniform and time-varying as cells grow, divide, and move. Single-GPU implementations remain constrained by compute capacity and memory, while prior reactive load-balancing approaches adjust partition boundaries only after imbalance is observed and do not exploit temporal trends in the workload. The authors develop a scalable multi-GPU framework for epidermal subcellular element tissue growth that combines GPU-parallel force evaluation, spatial binning, domain decomposition, workload-aware initial partitioning, and an RNN-guided load-balancing controller that learns residual corrections to a reactive boundary-adjustment rule.

Method

The authors propose a multi-GPU framework for accelerating epidermal tissue formation simulations based on a subcellular element model (SEM). The overall workflow encompasses workload-aware initialization, GPU-parallel element updates, ghost-region exchange, and RNN-guided runtime repartitioning.

In the SEM, each biological cell is represented by multiple interacting subcellular elements. The position $Y_{\alpha_i}$Yαi​​ of an element $\alpha_i$αi​ in cell $\alpha$α is updated based on forces from nearby elements and external structures:

where $\mathcal{I}$I is the set of neighboring elements, $V_{\mathrm{ele}}$Vele​ includes intra-cellular and inter-cellular interactions, and $V_{\mathrm{ext}}$Vext​ models external adhesion. All interactions are evaluated within a prescribed cutoff radius. To avoid the $O(n^2)$O(n2) cost of all-pairs comparisons, the 3D domain is organized into a regular grid of cubic bins sized by the cutoff distance. Each element only checks its own bin and the 26 neighboring bins, scaling the interaction calculation linearly. For multi-GPU execution, the spatial domain is divided into contiguous slices along one dimension and assigned to individual GPU ranks. Each rank maintains a ghost region containing elements within the cutoff distance of adjacent partitions to handle boundary interactions.

Because cell density and interaction counts vary spatially, a uniform geometric decomposition often leads to load imbalance. The authors compute a workload score for each cell $i$i:

where $n_i$ni​ is the number of subcellular elements in cell $i$i, and $n_{\mathcal{N}(i)}$nN(i)​ is the number of elements in the surrounding $3 \times 3 \times 3$3×3×3 bin neighborhood. The domain is divided into thin slices, and the workload of slice $s$s is $W_s = \sum_{i \in s} w_i$Ws​=∑i∈s​wi​. Contiguous slices are then grouped so that the accumulated workload for each GPU rank is approximately equal.

As cells grow and divide, the workload distribution changes dynamically. To address this, the authors introduce an RNN-guided dynamic load-balancing strategy that uses recent workload history and partition-state information to adapt the spatial decomposition.

The method starts with a reactive baseline that monitors computation times $T_t^-$Tt−​ and $T_t^+$Tt+​ of ranks adjacent to a partition boundary. A local imbalance $I_t = (T_t^+ - T_t^-) / ((T_t^+ + T_t^-) / 2)$It​=(Tt+​−Tt−​)/((Tt+​+Tt−​)/2) triggers a boundary shift when $|I_t| > \tau$∣It​∣>τ. The reactive adjustment magnitude is:

where $N_{\mathrm{slow}}$Nslow​ is the number of slices on the slower rank and $\alpha$α limits the adjustment.

To improve upon this reactive rule, a shared recurrent neural network (RNN) acts as a controller on partition boundaries. At each rebalancing step $t$t, a feature vector $\mathbf{x}_{b,t}$xb,t​ (including timing imbalance, partition widths, growth-cycle phase, and the reactive proposal) is fed into a gated recurrent unit (GRU) to update a hidden state $\mathbf{h}_{b,t}$hb,t​. The recurrent state is mapped to a scalar residual $r_{b,t}$rb,t​, which corrects the reactive proposal in a bounded step space:

This residual formulation allows the learned controller to refine the reactive rule while preserving a bounded adjustment magnitude. The corrected adjustments are applied jointly to all interior boundaries, shifting them toward the slower rank while respecting minimum partition width constraints.

The RNN controller is trained in a differentiable simulator of the load-balancing loop using synthetic workload traces. This allows the controller to learn the structure of the boundary-adjustment policy rather than specific biological dynamics. The input features are dimensionless local quantities, enabling generalization across different problem geometries and rank counts. The network is trained by backpropagation through time over rollouts of 50 rebalancing events. The training loss averages the relative imbalance of per-rank times and their normalized spread, with an $L_1$L1​ migration penalty added to discourage unnecessary boundary movement. The optimizer uses Adam with cosine annealing, and a soft relaxation is applied during training to keep the integer boundary assignment differentiable.

Experiment

The evaluation uses an epidermal subcellular element model on a GPU cluster to assess single-GPU acceleration, multi-GPU scalability, dynamic load balancing, and an RNN-based controller. GPU acceleration with spatial binning achieves near-linear scaling for large element counts, and multi-GPU execution benefits larger workloads. The RNN-guided load balancing reduces global timing imbalance with far fewer slice migrations than reactive or forecast-based methods, improving the balance–migration tradeoff. A developmental tissue growth simulation demonstrates the framework's capability for spatially evolving multicellular workloads.