Scalable Multi-GPU Simulation of 3D Multicellular Growth with RNN-Based Workload Balancing

Matvey Moisseyev; Huijing Du; Dandan Zheng; Chi Zhang; Hongfeng Yu

Source-linked AI summary

Scalable Multi-GPU Simulation of 3D Multicellular Growth with RNN-Based Workload Balancing

Matvey Moisseyev, Huijing Du, Dandan Zheng, Chi Zhang, Hongfeng Yu

arXiv:2608.25890v1Aug 26, 2026cs.DCcs.CEcs.LGmath.OC

TL;DR

Detailed SEM-based multicellular simulations face high element-interaction costs and workloads that change as tissues grow, divide, and move. The paper presents a multi-GPU framework with spatial binning, workload-aware partitioning, and an RNN-guided controller that learns history-based corrections to reactive repartitioning. GPU binning delivers roughly three orders of magnitude acceleration over a serial CPU baseline, while RNN-guided balancing reduces imbalance and migration in end-to-end simulations.

Problem
SEM simulations provide detailed cell-resolved mechanics but incur substantial element-interaction costs and dynamically changing spatial workloads that challenge single-GPU and static partitioning approaches.
Method
The framework combines GPU-parallel SEM interactions, spatial binning, multi-GPU domain decomposition, workload-aware initial partitioning, and an RNN controller learning residual corrections from workload and partition history.
Results
RNN-guided balancing reduces mean global imbalance from 11.3% under static partitioning to 3.5%, lowers total runtime by 9.0% relative to static partitioning, and migrates 7.7× fewer slices than the reactive baseline.
Takeaways & Limitations
History-aware load balancing improves workload balance while making more selective repartitioning decisions and reducing unnecessary slice migration in evolving multicellular simulations.

Abstract

from arXiv · show

Detailed multicellular growth simulations based on subcellular element models (SEMs) can capture complex tissue development, but their element-level interactions impose substantial computational cost. This work presents a scalable multi-GPU framework for 3D multicellular growth simulation that combines GPU acceleration, spatial binning, domain decomposition, and workload-aware partitioning. Cell movement, growth, and division continuously reshape the spatial workload distribution, causing initially balanced partitions to become inefficient over time. To address this, we introduce an RNN-based load-balancing controller that observes recent per-rank execution times and partition states and learns residual corrections to a reactive boundary-adjustment rule. The controller is trained offline in a differentiable surrogate of the load-balancing loop with randomized workload dynamics, requiring no measured execution traces for training. We evaluate the framework in terms of single-GPU acceleration, multi-GPU computation scaling, controller-level load-balancing behavior, and end-to-end simulation performance, with comparisons against static partitioning, reactive load balancing, and conventional time-series prediction baselines. A representative embryonic epidermal development use case further demonstrates the type of spatially and temporally evolving workload targeted by the framework. In our evaluation, GPU acceleration with spatial binning accelerates the interaction computation by roughly three orders of magnitude over a serial CPU baseline. RNN-guided load balancing reduces the mean global imbalance from 11.3% under static partitioning to 3.5%, lowers end-to-end runtime by 9.0% relative to static partitioning, and reduces slice migration by 7.7x compared with the reactive baseline, showing that history-aware control can improve workload balance while avoiding unnecessary repartitioning.

I. INTRODUCTION

Cell-resolved SEMs capture detailed tissue mechanics but create substantial computational demands as element counts and spatial distributions change. The work addresses this challenge with a multi-GPU framework and an RNN-guided load-balancing controller.

Cell-resolved SEMs represent each cell with multiple interacting elements, increasing force-evaluation costs while capturing fine-scale mechanical interactions.
A single GPU is limited by computational capacity and device memory, motivating spatial decomposition and boundary-data exchange across multiple GPUs.
Cell growth, division, and movement continuously change local element density, so initially balanced partitions can become imbalanced during execution.
The framework combines GPU-parallel force evaluation, spatial binning, distributed domain decomposition, ghost regions, and workload-aware initial partitioning.
The RNN controller uses recent workload and partition-state history to learn residual corrections to a reactive boundary-adjustment rule.
The framework is evaluated through single-GPU, multi-GPU, controller-level, and end-to-end simulation experiments.

B. Parallel and GPU-Accelerated Multicellular Simulation

Prior multicellular simulation work establishes GPU acceleration and distributed execution as routes to larger models, while evolving workloads require adaptive partitioning. This framework applies those ideas to epidermal SEM simulation with RNN-guided runtime repartitioning.

Existing systems provide parallel infrastructure for cell-based modeling, including general frameworks and large three-dimensional population simulators.
GPU methods accelerate multicellular mechanics by evaluating many local interaction calculations concurrently, including three-dimensional epidermal SEMs.
Dynamic particle distributions motivate adaptive domain decomposition because computation assigned to partitions can change substantially during execution.
Predictive load balancing uses workload forecasts or temporal models to guide partition adjustment and migration decisions.
In SEM tissue growth, proliferation and local density changes create spatially localized imbalance affecting GPU computation and inter-GPU boundary exchange.
The workflow combines workload-aware initialization, GPU element updates, ghost-region exchange, and RNN-guided runtime repartitioning.

B. Computational Characteristics

SEM interaction costs are high because dense, evolving element populations produce nonuniform workloads. Spatial binning and neighboring-rank ghost regions support efficient local interaction evaluation and distributed execution.

A naive all-pairs implementation requires O(n^2) candidate interaction tests for n elements.
Dense regions require more interaction evaluations than sparse regions, while movement, growth, and division continually alter element distributions.
The resulting requirements are efficient local interaction evaluation and scalable distribution of dynamically changing spatial workloads across GPUs.
Spatial binning restricts each element’s candidate search to its own bin and the 26 neighboring bins.
Under bounded local density, binning keeps candidate interactions per element approximately constant and makes interaction calculation scale linearly with element count.
Contiguous spatial slices are assigned to GPU ranks, with ghost regions exchanging cutoff-distance boundary information only between neighbors.

E. Workload-Aware Initial Partitioning

The workload-aware initial partitioning method estimates computational cost from element counts and local neighborhood density, then groups contiguous slices so GPU ranks receive approximately equal workloads. Its balance reflects only the simulation’s initial state.

Uniform geometric decomposition can assign unequal computational workloads when equal-sized subdomains contain different element counts or interaction densities.
The workload score accounts for each cell’s subcellular-element count and the number of elements in its surrounding 3 × 3 × 3 bin neighborhood.
The domain is divided along one axis into thin slices, each receiving a workload value that forms a one-dimensional histogram.
Figure 2 contrasts nonuniform slice workloads, unequal uniform-geometric assignments, and workload-aware boundaries across 8 GPU ranks.
Contiguous slices are grouped so each GPU rank receives approximately equal accumulated workload.
The initial strategy improves balance for spatially heterogeneous workloads but cannot prevent runtime imbalance as the tissue evolves.

IV. PREDICTIVE DYNAMIC LOAD BALANCING

The method combines reactive boundary adjustment with history-aware RNN residual corrections to anticipate evolving imbalance and redistribute slices between neighboring GPU ranks.

Predictive Dynamic Load Balancing: Temporal workload correlations allow recent workload and partition states to inform subsequent imbalance changes.This motivates predictive adjustment rather than relying only on performance degradation already observed.
Predictive Dynamic Load Balancing: The RNN maintains a hidden state per partition boundary and adds a residual correction to the reactive boundary proposal.The corrected movement shifts the boundary and migrates slices between neighboring GPU ranks.
Reactive Baseline: Reactive balancing monitors neighboring-rank execution times and shifts boundaries after observed imbalance exceeds a threshold.The slower rank transfers workload to the faster rank, with adjustment magnitude controlled by α.

C. RNN-Based Boundary Control

The controller is a shared recurrent boundary model that uses local timing, workload, and partition-state features to produce bounded residual corrections adaptable across GPU-rank counts.

C. RNN-Based Boundary Control: A shared recurrent cell maintains a separate hidden state for each interior partition boundary.Shared recurrent weights and locally defined inputs allow application to different numbers of GPU ranks.
C. RNN-Based Boundary Control: Boundary features include local and global imbalance, partition widths and slack, prior movement, growth-cycle state, recent growth events, and the reactive proposal.All features are dimensionless ratios clamped to fixed ranges for deployment.
C. RNN-Based Boundary Control: The recurrent state maps to a scalar residual that corrects the reactive proposal within a maximum boundary-movement bound.Clipping and the bounded residual formulation preserve the adjustment limit while refining the reactive rule.
C. RNN-Based Boundary Control: A single-layer GRU with hidden-state size 24 and approximately 2.8 × 10^3 trainable parameters implements the controller.Zero initialization of the output layer makes the untrained controller reproduce the reactive proposal.

D. Residual Boundary Adjustment

Residual adjustments are applied jointly under partition-width constraints, while offline differentiable training uses randomized workload dynamics and migration-aware objectives.

D. Residual Boundary Adjustment: All interior boundaries are updated in one sweep toward slower neighboring ranks, with movements truncated to preserve minimum partition widths.The updated boundaries define the next simulation subdomains before slice migration.
D. Residual Boundary Adjustment: The residual formulation separates reactive direction and feasibility constraints from the RNN’s learned magnitude correction.Boundary-update rules enforce admissible partition sizes before migration.
D. Residual Boundary Adjustment: The controller trains in a differentiable load-balancing simulator with randomized workload evolution, growth events, bursts, measurement jitter, and noise-free objective timing.Synthetic traces target boundary-policy structure rather than detailed biological dynamics.
D. Residual Boundary Adjustment: Randomized rank counts, domain sizes, and minimum widths, with unseen evaluation seeds, support generalization across problem geometries.Dimensionless features and shared per-boundary weights are part of this design.
D. Residual Boundary Adjustment: Training uses 50-event backpropagation-through-time rollouts, while evaluation includes 100-event unseen scenarios and a 400-event stability rollout.The objective combines imbalance and normalized spread with an L1 penalty of 5 × 10^-4 per moved slice.
D. Residual Boundary Adjustment: The RNN is used only for dynamic rebalancing; initial decomposition still uses computed workload scores.Training also uses soft differentiable boundary assignment with annealed temperature.

V. EXPERIMENTAL EVALUATION

The evaluation measures GPU acceleration, multi-GPU scalability, dynamic balancing, and controller comparisons using an epidermal SEM on the Swan cluster.

V. EXPERIMENTAL EVALUATION: The experiments assess acceleration, multi-GPU scaling, dynamic load-balancing effectiveness, and RNN comparisons with static, reactive, and time-series baselines.A representative multicellular growth use case demonstrates the targeted scientific capability.
Experimental Setup: Experiments use epidermal SEM parameters on NVIDIA L40S GPUs with one simulation rank assigned per GPU unless otherwise stated.The cluster uses MPI communication over HDR100 InfiniBand.
Single-GPU Acceleration: GPU acceleration experiments compare serial CPU all-pairs, GPU all-pairs, and GPU spatial-binning implementations over increasing element counts.Figure 4 reports interaction-evaluation time over 1,000 iterations on a log–log scale.
Single-GPU Acceleration: At 1.6 × 10^5 elements, the binned GPU completes in 7.2 s versus roughly 2 hours for serial CPU, a speedup of roughly three orders of magnitude.The GPU all-pairs implementation completes in 14.1 s at the same size.
Single-GPU Acceleration: At the largest configuration, spatial binning takes 255 s versus 2,355 s for GPU all-pairs, yielding a 9.2× speedup.Increasing elements approximately 40× raises binned runtime 35×; only inputs below 10^3 elements make the implementations comparable.

C. Multi-GPU Scalability

Multi-GPU execution reduces core GPU computation time, but fixed-size strong scaling is sublinear and saturates when per-GPU workloads become too small. RNN-guided balancing achieves the lowest imbalance while reducing runtime and slice migration relative to the evaluated alternatives.

Strong scaling: GPU computation time decreases as additional GPUs are used, demonstrating benefits from distributed execution for the core interaction computation.Figure 5 isolates GPU computation by excluding startup, load-balancing, communication, and I/O costs.
Strong scaling: At 8 GPUs, computation-only parallel efficiency is approximately 45% for 435K cells versus 36% for 60K cells.The 60K-cell case begins saturating by 8 GPUs, while the larger case provides more computation per rank.
Strong scaling: The 435K-cell case benefits from additional GPUs up to 16, but efficiency drops to 31% as the fixed workload is divided across more devices.Even the larger case does not fully saturate all GPUs at 16 devices.
Dynamic load balancing: Static partitioning produces persistent per-rank timing differences, while reactive and RNN-guided balancing reduce the spread during the simulation.The RNN-guided method further tightens computation-time curves and slightly lowers per-rank computation times.
Dynamic load balancing: The RNN-guided method achieves the lowest mean global imbalance and reduces total runtime relative to static partitioning.Its runtime is comparable to, and slightly lower than, the reactive baseline in this experiment.
Dynamic load balancing: RNN-guided balancing maintains balanced GPU workloads while substantially reducing data movement required during repartitioning.Compared with reactive balancing, it migrates far fewer slices through more selective boundary movement.

E. RNN Controller Evaluation

The controller-level evaluation compares static, forecast-based, reactive, and RNN policies on held-out scenarios using time-averaged imbalance and total slice migration. The RNN achieves a conservative balance–migration tradeoff, especially at high rank counts, while controller-level gains require end-to-end validation.

Evaluation setup: The evaluation compares static partitioning, persistence and Holt forecast-then-repartition, reactive balancing, and the RNN controller on held-out scenarios.All methods use hard integer partition boundaries and are assessed on the same scenarios.
Evaluation metrics: Time-averaged global imbalance and total migrated slices measure the balance–migration tradeoff across each rollout.Migration counts slices moved over the full rollout rather than average migration per step.
Results: At P = 16, persistence and Holt migrate 6481 and 5243 slices, respectively, while producing higher imbalance than static partitioning.Thin slices make frequent forecast-driven boundary movement less effective at this rank count.
Results: At P = 16, reactive and RNN controllers keep imbalance at or below the static level using fewer than 60 slice migrations.At P = 2, static partitioning is already balanced within 0.5%, so additional movement can be counterproductive.
Results: Across rank counts, the RNN migrates the fewest slices among adaptive methods while keeping time-averaged imbalance below the static baseline.It does not always achieve the lowest imbalance at moderate rank counts, but uses the most conservative migration schedule.
End-to-end relevance: Controller-level gains alone do not establish end-to-end usefulness, so the framework is also evaluated through full multi-GPU simulations.The paper separately examines whether isolated imbalance reductions translate into improved simulation performance.

F. Scientific Use Case: Embryonic Epidermal Development

Embryonic epidermal development serves as a representative use case for simulations whose cell populations and workloads evolve spatially and temporally. The framework combines scalable multi-GPU execution with history-aware balancing to address these dynamics, while larger-scale validation remains future work.

Scientific use case: Embryonic epidermal development is used as a representative scientific use case for scalable multicellular growth simulation.The simulation covers proliferation, migration, and differentiation during tissue stratification.
Framework: The framework combines GPU acceleration, spatial binning, domain decomposition, and workload-aware initial partitioning for three-dimensional multicellular growth simulation.These components target the cost of element-level interactions and nonuniform workloads in growing tissues.
Dynamic balancing: RNN-guided balancing uses recent workload and partition-state history to make more selective repartitioning decisions as cell movement, growth, and division alter workload distribution.The controller learns residual corrections to a reactive boundary-adjustment rule.
End-to-end results: In 8-GPU 435K-cell simulations, RNN-guided balancing reduces mean global imbalance from 11.3% under static partitioning to 3.5% and total runtime by 9.0%.It also migrates 7.7× fewer slices than the reactive baseline while maintaining slightly better load balance and comparable runtime.
Workload evolution: The use case produces spatially and temporally nonuniform workloads through proliferation, differentiation, and tissue reorganization.Snapshots span Day 0 to Day 5 and illustrate the evolving tissue population.
Scope and future work: The evaluation plans to extend to million-cell-scale models because fixed-size strong-scaling experiments may underutilize GPUs at high device counts.Larger simulations are expected to provide more computation per GPU and further test scalability.