PRISM: Profiling-Free Symbolic Memory-Driven Strategy Planner for Large DNN Model Training

Read full text
Authors: Ruiwen Wang, Philippe Fang, Chong Li, Thibaut Tachon, Raja Appuswamy
Cited by: 0
References: 6

Abstract

The rapid growth of large-scale deep neural networks (DNNs) has introduced severe memory and performance bottlenecks during distributed training. Existing automated planners for parallelization strategies often rely heavily on profiling or empirical tuning, which significantly increases engineering cost and wastes large-scale cluster resources. In this work, we present PRISM, a profiling-free, symbolic memory-driven strategy planner for large DNN training. PRISM introduces a unified symbolic memory cost model that captures the layered structure of modern architectures and integrates with a communication model to evaluate trade-offs across data, tensor, pipeline, virtual pipeline, expert, and sequence parallelism, as well as activation recomputation and optimizer sharding. By formulating strategy selection as an optimization problem, PRISM identifies globally optimal parallel strategies under device memory budgets. Our evaluation across representative large models demonstrates that PRISM achieves accurate memory prediction and substantial improvements in Model FLOPs Utilization (MFU), reducing bubble and communication overheads without costly profiling.

Fields of Study & Subjects

Advanced Neural Network Applications Machine Learning in Materials Science Parallel Computing and Optimization Techniques

Publication details

Pages
308–319
Publication date
9 January 2026
Language
EN