AI/ ai · hardware · open-source · research

STAGE Lets Researchers Simulate LLM Training at 128K-GPU Scale

A new open-source framework synthesizes high-fidelity execution graphs for distributed AI workloads, no cloud supercluster required.

Researchers have released STAGE, a tool that models how large language models behave across massive distributed hardware setups — without needing access to the hardware itself.

STAGE, short for Symbolic Tensor Graphs for distributed AI workload synthesis, generates detailed execution graphs that simulate LLM and mixture-of-experts training and inference runs. The framework supports a range of parallelization strategies and can model configurations spanning more than 128,000 GPUs while preserving accuracy at the individual tensor level across compute, memory, and communication dimensions. The code is publicly available on GitHub under the astra-sim project.

This matters because studying how to optimize AI systems at scale — tuning parallelization, exploring hardware designs before silicon is taped out — currently requires either owning a supercluster or having a cloud provider hand you access to one. That shuts most academic researchers out entirely. STAGE offers a credible alternative: synthesize the traces, run the analysis, and never touch a real A100.

The broader trend here is that AI infrastructure research is quietly bifurcating into haves and have-nots. Tools like STAGE won't close that gap entirely, but they push back against the assumption that serious systems work can only happen inside a hyperscaler's data center.

TR

The Revision

Written by an AI system from the public sources credited above. How we write →