Scalable Synthesis of distributed LLM workloads through Symbolic Tensor Graphs

Changhai Man; Joongun Park; Hanjiang Wu; Huan Xu; Srinivas Sridharan; Tushar Krishna

cs.DCNov 13, 2025

Scalable Synthesis of distributed LLM workloads through Symbolic Tensor Graphs

Authors: Changhai Man, Joongun Park, Hanjiang Wu, Huan Xu, Srinivas Sridharan, Tushar Krishna

Organizations: Georgia Institute of Technology · NVIDIA Inc.

Abstract

Optimizing the performance of large language models (LLMs) on large-scale AI training and inference systems requires a scalable and expressive mechanism to model distributed workload execution. Such modeling is essential for pre-deployment system-level optimizations (e.g., parallelization strategies) and hardware design-space explorations. While recent efforts have proposed collecting execution traces from real systems, access to large-scale infrastructure remains limited to major cloud providers. Moreover, traces capturing execution on a specific platform cannot be easily adapted to study alternate software and/or hardware configurations, especially at scale. We introduce STAGE, a framework that synthesizes high-fidelity execution graphs to accurately model distributed AI workloads (including LLMs and MoEs). STAGE supports a comprehensive set of parallelization strategies, allowing users to systematically explore a wide spectrum of model architectures and system configurations. STAGE demonstrates its scalability by synthesizing high-fidelity LLM traces spanning over 128K GPUs, while preserving tensorlevel accuracy in compute, memory, and communication. STAGE is publicy available at https://github.com/astra-sim/stage

Explore similar work

A Few GPUs, A Whole Lotta Scale: Faithful LLM Training Emulation with PrismLLM
May 15, 2026Shaoke Xi, ChonLam Lao, Boyi Jia +11Large Models Graphics Processing Unit Memory
AAFLOW: Scalable Patterns for Agentic AI Workflows
May 4, 2026Arup Kumar Sarker, Mills Staylor, Aymen Alsaadi +3Agentic Workflows Modular Agentic Framework
Scepsy: Serving Agentic Workflows Using Aggregate LLM Pipelines
Apr 16, 2026Marcel Wagenländer, Otto White, Britannio Jarrett +6Agentic Workflows Llm-Based Agent

May 15, 2026 · cs.DCJ/K move · Enter open · S save

A Few GPUs, A Whole Lotta Scale: Faithful LLM Training Emulation with PrismLLM

Shaoke Xi, ChonLam Lao, Boyi Jia, Jiaqi Gao+10

Alibaba Group · Harvard University · Shanghai Jiao Tong University · Zhejiang University

Large language model (LLM) training today runs on clusters spanning thousands of GPUs. While this scale enables rapid model advances, developing, debugging, and performance-tuning the training framework inevitably becomes complex and costly. This is because engineers often need to reproduce production behaviors to diagnose failures or evaluate optimizations, thereby demanding frequent and even exclusive access to production-scale clusters -- which becomes increasingly hard given that the majority of GPUs are already committed to production workloads. Simulation relies on complex performance models that are difficult to maintain, and downscaled experiments often fail to capture scale-dependent behaviors. We present PrismLLM to decouple large-scale execution from the need to access large clusters, enabling engineers to run and observe ranks of interest under faithful large-scale behavior using only a few GPUs. PrismLLM constructs a high-fidelity execution graph via a slicing-based approach that captures computation, communication, and dependencies of the target scale. Then, PrismLLM performs hybrid emulation where selected ranks execute the original program while the remaining ranks are replayed as virtual participants. Experiments on large-scale LLM training workloads show that PrismLLM accurately reproduces performance and memory behavior, achieving only 0.58% average error in iteration time and less than 0.01% error in peak GPU memory usage. PrismLLM can emulate clusters of up to 8192 GPUs using fewer than 1% of the physical GPUs required by the original deployment.

Large Models Graphics Processing Unit Memory Emulators

Open paper Ask about it PDF

Abstract

Explore similar work

A Few GPUs, A Whole Lotta Scale: Faithful LLM Training Emulation with PrismLLM

AAFLOW: Scalable Patterns for Agentic AI Workflows

Scepsy: Serving Agentic Workflows Using Aggregate LLM Pipelines