Scalable Multi-GPU Simulation of 3D Multicellular Growth with RNN-Based Workload Balancing
Authors: Matvey Moisseyev, Huijing Du, Dandan Zheng, Chi Zhang, Hongfeng Yu
Organizations: School of Computing, University of Nebraska–Lincoln · Department of Mathematics, University of Nebraska–Lincoln · Department of Radiation Oncology, University of Rochester Medical Center · School of Biological Sciences, University of Nebraska–Lincoln · Holland Computing Center, University of Nebraska–Lincoln
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.
Figures & tables
Fig. 1: Overview of the proposed multi-GPU multicellular growth simulation framework. The framework initializes the simulation with workload-aware partitioning, distributes subdomains across GPU ranks, and iteratively performs spatial binning, distributed GPU element updates, ghost-region exchange, and cell growth and division. At each rebalancing interval, timing and partition-state information are used by the RNN-guided controller to update partition boundaries and migrate slices. The final tissue state is assembled from the local subdomains.
Fig. 2: Schematic illustration of workload-aware initial partitioning. (a) The number of subcellular elements varies across slices along the partitioning axis, producing a spatially nonuniform workload. (b) Uniform geometric partitioning assigns equal-width slice ranges to GPU ranks, which can produce unequal workload shares when the workload is spatially nonuniform. (c) The workload-aware method selects partition boundaries using the accumulated workload score Ws so that each rank receives approximately the same total workload. This example uses 8 GPU ranks, labeled 0–7. Percentages are rounded workload shares assigned to each GPU rank. Practical workload profiles are often more irregular.
Fig. 3: RNN-guided dynamic load balancing. For each interior partition boundary, timing, workload-history, and partition-state features are collected and passed to both a reactive boundary proposal and a shared recurrent controller. The RNN outputs a residual correction that is combined with the reactive proposal to produce a bounded boundary adjustment. The corrected adjustment shifts the partition boundary and migrates slices between neighboring GPU ranks.
Fig. 4: Single-GPU performance of element-level interaction evaluation measured over 1,000 simulation iterations on a log–log scale as the number of subcellular elements increases. The figure compares the original serial CPU all-pairs implementation, the GPU all-pairs implementation, and the GPU implementation with spatial binning. Reference slopes indicate the expected O(n2) and O(n) scaling trends.
Fig. 5: Multi-GPU strong-scaling performance of GPU computation time for two multicellular simulation sizes. (a) GPU computation time as the number of GPUs increases. (b) Corresponding speedup, with the dashed line indicating ideal linear scaling.
Fig. 6: Evolution of per-rank computation time during the simulation for static partitioning, reactive load balancing, and RNN-guided load balancing over a 100,000 -iteration simulation. Each curve represents one GPU rank, and each point reports the rank’s computation time over one simulation interval. The x-axis reports simulation iterations in units of 103 .
Fig. 7: End-to-end performance of static partitioning, reactive load balancing, and RNN-guided load balancing over a 100,000 -iteration simulation. The figure compares total runtime, mean global imbalance, and slice migration cost. Dynamic load balancing reduces timing imbalance, while slice migration measures the amount of data movement introduced by repartitioning.
Fig. 8: Controller-level evaluation on held-out scenarios. (a) Time-averaged global imbalance for static partitioning, persistence-based forecast-then-repartition, Holt-based forecast-then-repartition, the reactive baseline, and the RNN controller. (b) Migration cost measured as the total number of slices migrated over the rollout. Each point corresponds to one held-out rollout on the 435K-cell initial condition, with imbalance averaged over the measured intervals and migration accumulated over the rollout. The vertical axis is logarithmic in both panels, with a linear region below 10 in (b) so that static partitioning, which migrates no slices, remains visible. Lower values are better for both metrics.
Fig. 9: Representative epidermal development simulation at multiple time points. Cells are rendered as colored ellipsoids, with different colors indicating different cell types or states. The snapshots show the evolving cell population over a 5-day developmental period, illustrating progressive proliferation, differentiation, and tissue reorganization. The resulting spatially and temporally nonuniform workload motivates the proposed multi-GPU framework and dynamic load-balancing strategy.
Deep learning training commonly runs on shared multi-tenant GPU servers, where exclusive allocation provides isolation but can leave resources underutilized and increase queueing time. Collocation can improve efficiency, but interference-agnostic placement may cause severe slowdowns, while inaccurate memory information can lead to out-of-memory (OOM) failures. We present AEGIS, a server-scale runtime scheduling system for controlled collocation of deep learning training workloads on shared multi-GPU servers. AEGIS integrates memory feasibility, post-placement observation, runtime-pressure filtering, placement, and OOM-aware recovery in a single scheduling loop. After placement, AEGIS observes workload activity before permitting further collocation, then uses low-overhead telemetry to determine whether a GPU can safely accept additional work. OOM failures trigger retries under progressively safer memory conditions, eventually falling back to exclusive execution. This online approach avoids costly offline pairwise compatibility profiling. We evaluate AEGIS using vision, Transformer, recommendation, and LLM-style workloads across three production-derived traces. AEGIS reduces geometric-mean makespan by 16% relative to Lucid, 21% relative to Horus, and 27% relative to exclusive allocation. Sensitivity studies show that activity-anchored observation and runtime-pressure filtering balance conservative isolation against interference-agnostic collocation, improving makespan while limiting sharing-induced per-task slowdown.
Ehsan Yousefzadeh-Asl-Miandoab, Büşra Karatay Demiray, Florina M. Ciorba +2
IT University of Copenhagen Denmark · University of Basel Switzerland
The rapid growth of large-scale machine learning (ML) has made distributed training across multiple GPUs a fundamental component of modern ML systems. As model sizes and computational throughput continue to increase, communication overhead has become a dominant bottleneck in multi-GPU training, particularly when computation and communication are executed sequentially. This work explores concurrent execution of computation and collective communication using two portable runtime controls: shared-memory-driven occupancy shaping for computation kernels and elevated scheduling priority for communication kernels. Our approach regulates computation-kernel residency through per-block shared-memory allocation, leaving sufficient on-chip resources for communication kernels to make progress. In addition, assigning higher priority to communication streams ensures steady communication progress once resources become available. Experiments on NVIDIA A40, A100, H100, and AMD MI250X GPUs demonstrate that the proposed method enables effective computation-communication overlap and reduces total execution time by up to 25.5 percent, without modifying vendor libraries or kernel implementations.
Minyu Cui, Miquel Pericas
Chalmers University of Technology and University of Gothenburg, Gothenburg, Sweden
Simulation-based RL for contemporary robot control is increasingly organized around GPU-resident simulation: physics, rollout collection, and learning are placed on a single GPU-centric execution path. This paradigm has greatly improved training speed, but it has also encouraged a default assumption that efficient training requires physics to reside on the GPU. We revisit this assumption. Our view is that, in simulation-dominated robot control, the essential question is not which processor runs physics, but whether simulation throughput, policy learning, and runtime synchronization form an efficient end-to-end loop. We present UniLab, a heterogeneous CPU-simulation / GPU-learning architecture that decouples CPU-parallel simulation from GPU policy updates through a unified runtime for data movement, buffering, and synchronization. UniLab is implemented as a complete and extensible training system using MuJoCoUni and MotrixSim CPU-batched physics backends, supporting PPO, FastSAC, FlashSAC, and APPO. On representative simulation-based robot control tasks, UniLab improves end-to-end training efficiency by 3--10× under the same hardware configuration, while reducing dependence on the NVIDIA CUDA-based software stack and supporting cross-platform execution on the Apple macOS platform and the AMD ROCm and Intel XPU accelerator backends. These results show that GPU simulation is an effective path to efficient training, but not a necessary one, broadening the practical system choices available for robot RL training. Project page: https://unilabsim.github.io.