Maya: Optimizing Deep Learning Training Workloads using GPU Runtime Emulation

TL;DR

Maya tackles the challenge of optimizing deep learning training at scale by removing the semantic gap between workloads and their performance representations. It does so with transparent device emulation that intercepts accelerator APIs from unmodified training code, coupled with a trace-driven end-to-end simulator and pluggable kernel estimators. The approach delivers high-fidelity end-to-end predictions (often <5% error) and enables rapid configuration search that achieves near-optimal training costs (up to 56% savings compared to baselines). Its architecture supports hyperscale workloads and generalizes across frameworks and model types, reducing the barrier to exploring and deploying efficient training recipes. Overall, Maya provides a practical, deployment-free path to scalable, cost-efficient DL training optimization.

Abstract

Training large foundation models costs hundreds of millions of dollars, making deployment optimization critical. Current approaches require machine learning engineers to manually craft training recipes through error-prone trial-and-error on expensive compute clusters. To enable efficient exploration of training configurations, researchers have developed performance modeling systems. However, these systems force users to translate their workloads into custom specification languages, introducing a fundamental semantic gap between the actual workload and its representation. This gap creates an inherent tradeoff: systems must either support a narrow set of workloads to maintain usability, require complex specifications that limit practical adoption, or compromise prediction accuracy with simplified performance models. We present Maya, a performance modeling system that eliminates these tradeoffs through transparent device emulation. By operating at the narrow interface between training frameworks and accelerator devices, Maya can capture complete workload behavior without requiring code modifications or translations. Maya intercepts device API calls from unmodified training code to directly observe low-level operations, enabling accurate performance prediction while maintaining both ease of use and generality. Our evaluation shows Maya achieves less than 5% prediction error across diverse models and optimization strategies, identifying configurations that reduce training costs by up to 56% compared to existing approaches.

Figures (16)

• Figure 1: Existing deep learning training performance modeling systems struggle with a tradeoff between fidelity, usability, and generality. Naive analytical models lack fidelity, while the advanced modeling approaches make a tradeoff between usability and generality. Maya breaks this tradeoff through a novel device emulation approach, achieving all three simultaneously.
• Figure 2: Sensitivity of optimal training configurations to cluster size for GPT-3 18.4B on H100 GPUs. As GPU counts increase, configurations shift fundamentally --- from memory-efficient combinations of tensor and pipeline parallelism in smaller clusters to higher data-parallel degrees in larger clusters. The cross-deployment cost matrix highlights that deploying configurations tuned for one cluster size on another can lead to inefficiencies, increasing costs by up to 74% due to suboptimal resource use. These results emphasize the necessity of scenario-specific configuration tuning, a key challenge Maya addresses through precise performance modeling.
• Figure 3: The user workflow across three systems. With AMPed (top), the user provides a declarative configuration specifying high-level parameters which are then fed into a predefined analytical model in the system. If a new model architecture or performance optimization is introduced, the system is rendered unusable. In Proteus (middle), the user must manually translate their entire model into a custom format and write a separate "Strategy Tree" to explicitly define the parallelization strategy. With Maya (bottom), the user runs their original, unmodified training script, and the system automatically captures a low-level execution trace through transparent emulation, requiring no user intervention.
• Figure 4: Comparison of modeling approaches: Traditional systems require explicit workload specification and several complex heuristics steps to obtain a kernel level execution graph apt for simulation. On the other hand, Maya directly captures the computation graph at a lower-level through transparent device emulation.
• Figure 5: Maya architecture: (1) Unmodified training code is executed through a virtual runtime that emulates the device drivers given emulation specs, (2) Worker traces capturing device API calls are merged into a unified trace, (3) Kernels in the unified trace are annotated with predicted runtimes using pre-trained estimators, (4) An event-driven simulator processes the annotated traces using cluster specifications to produce performance predictions.