Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

STAGE: A Symbolic Tensor grAph GEnerator for distributed AI system co-design

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

TL;DR

STAGE addresses the scarcity and narrow applicability of real execution traces for distributed AI workloads by introducing Symbolic Tensor Graph Generator, a framework that synthesizes high-fidelity execution graphs using a symbolic tensor representation. It supports diverse LLM architectures and parallelism strategies, including tensor- and graph-level distributions, and outputs Chakra-based execution graphs that encode compute, memory, and communication dependencies. The approach is validated against real traces and demonstrated at scale up to 32K GPUs, while remaining adaptable to future system configurations and architectures. Publicly available and scalable, STAGE enables fast, systematic design-space exploration and system-level optimization for large-scale AI infrastructure.

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 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 obtained from existing platforms cannot be easily adapted to study future larger-scale system configurations. We introduce Symbolic Tensor grAph GEnerator(STAGE), a framework that synthesizes high-fidelity execution traces to accurately model LLM workloads. STAGE supports a comprehensive set of parallelization strategies, allowing users to systematically explore a wide spectrum of LLM architectures and system configurations. STAGE demonstrates its scalability by synthesizing high-fidelity LLM traces spanning over 32K GPUs, while preserving tensor-level accuracy in compute, memory, and communication. STAGE is publicly available to facilitate further research in distributed machine learning systems: https://github.com/astra-sim/symbolic tensor graph

STAGE: A Symbolic Tensor grAph GEnerator for distributed AI system co-design

TL;DR

STAGE addresses the scarcity and narrow applicability of real execution traces for distributed AI workloads by introducing Symbolic Tensor Graph Generator, a framework that synthesizes high-fidelity execution graphs using a symbolic tensor representation. It supports diverse LLM architectures and parallelism strategies, including tensor- and graph-level distributions, and outputs Chakra-based execution graphs that encode compute, memory, and communication dependencies. The approach is validated against real traces and demonstrated at scale up to 32K GPUs, while remaining adaptable to future system configurations and architectures. Publicly available and scalable, STAGE enables fast, systematic design-space exploration and system-level optimization for large-scale AI infrastructure.

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 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 obtained from existing platforms cannot be easily adapted to study future larger-scale system configurations. We introduce Symbolic Tensor grAph GEnerator(STAGE), a framework that synthesizes high-fidelity execution traces to accurately model LLM workloads. STAGE supports a comprehensive set of parallelization strategies, allowing users to systematically explore a wide spectrum of LLM architectures and system configurations. STAGE demonstrates its scalability by synthesizing high-fidelity LLM traces spanning over 32K GPUs, while preserving tensor-level accuracy in compute, memory, and communication. STAGE is publicly available to facilitate further research in distributed machine learning systems: https://github.com/astra-sim/symbolic tensor graph

Paper Structure

This paper contains 36 sections, 4 equations, 13 figures, 13 tables.

Table of Contents

  1. Introduction
  2. Background and Related Works
  3. Large Language Models (LLMs)
  4. Distributed Training with Multi-dimensional Parallelization Strategies
  5. Execution Traces
  6. Motivation
  7. Challenge 1: Limitations of Real-System ETs
  8. Challenge 2: Modeling LLM Architecture and Parallelism Diversity
  9. STAGE: Symbolic Tensor Graph Generator
  10. STAGE Overview
  11. Workload definition
  12. Input: Model and Module Templates
  13. Output: Execution Graph.
  14. Symbolic Tensor Representation
  15. Workload Distributor
  16. ...and 21 more sections

Figures (13)

  • Figure 1: Overview of STAGE
  • Figure 2: An Example of Execution Trace for GPU Operations.
  • Figure 3: STAGE Generation Flow Overview.
  • Figure 4: Using Symbolic Tensor Representation to Annotate MultiHead Attention with Sequence and Tensor Parallelism.
  • Figure 5: Tensor Distribution Mismatch: After applying tensor-level distribution
  • ...and 8 more figures