Pith. sign in

REVIEW 4 major objections 7 minor 79 references

STAGE generates realistic execution traces for distributed LLM training from a compact symbolic tensor graph, matching measured cluster behavior closely enough to serve as a stand-in for expensive real runs.

Reviewed by Pith at T0; open to challenge. T0 means a machine referee read the full paper against a public rubric. the ladder, T0–T4 →

T0 review · deepseek-v4-flash

2026-08-03 22:24 UTC pith:X2FMTQ5E

load-bearing objection A solid, well-scoped workload-generation framework with a genuinely new symbolic-tensor IR and Pull/Push collective matcher; the main weakness is that compute-time 'fidelity' is only validated in-sample on the benchmarked H100 stack. the 4 major comments →

arxiv 2511.10480 v3 pith:X2FMTQ5E submitted 2025-11-13 cs.DC cs.AI

Scalable Synthesis of distributed LLM workloads through Symbolic Tensor Graphs

classification cs.DC cs.AI
keywords symbolic tensor graphexecution trace synthesisdistributed LLM trainingparallelism strategiesmixture-of-expertsdesign space explorationworkload modelingexecution graph
verification ladder T0 review T1 audit T2 compute T3 formal T4 reserved

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

STAGE aims to make distributed-LLM workload modeling cheap and configurable by replacing real execution-trace collection with synthesis. The paper claims that a symbolic tensor graph—an intermediate representation in which tensor shapes carry distribution annotations like duplicate, partition, or partial-sum—can encode the structure of dense, mixture-of-experts, and state-space models, along with any combination of data, tensor, sequence, pipeline, and expert parallelism. From that graph, STAGE instantiates a full execution graph with per-tensor compute, memory, and communication, and validates it against traces from a 128-GPU H100 cluster: memory errors of a few percent, communication-volume errors below about 3%, and compute-time errors averaging 4.25%. If correct, system designers can explore parallelization strategies and hypothetical hardware without owning a multi-thousand-GPU cluster.

Core claim

The central claim is that a graph of symbolic tensor operations, annotated with sharding semantics, is a sufficient generative model for distributed LLM execution traces. STAGE assembles module templates (attention variants, feed-forward networks, MoE layers, normalization) into a whole model, distributes the workload across devices by rewriting tensor shapes with parallel dimensions, detects sharding mismatches between producers and consumers, inserts the corresponding collective communications, and then converts the symbolic graph to a numeric execution graph. The paper reports that this workflow reproduces operator counts, communication volumes, and peak memory within a few percent of mea

What carries the argument

The Symbolic Tensor Graph (STG), an intermediate representation where tensor shapes include symbolic dimensions and distribution markers (Duplicated, Partition, PartialSum), with all multiplications expressed as einsum. The load-bearing mechanism is the producer-consumer communication matcher: it views any sharding change as a Pull to a virtual node followed by a Push to consumers, which systematically derives the required collective (AllGather, AllReduce, ReduceScatter, AllToAll, or combinations) whenever a producer's output layout differs from what a consumer expects. This enables arbitrary hybrid parallel strategies without hand-coding each one.

Load-bearing premise

The compute-time fidelity rests on a lookup table and a roofline model calibrated on the same H100 cluster used for validation, so the reported accuracy is in-sample; if those calibration constants change on other GPUs, networks, or software stacks, the claimed tensor-level accuracy at 32K GPUs would not transfer automatically.

What would settle it

Take a standard dense LLM (say, a 70B-parameter transformer), collect fresh execution traces on a different accelerator generation without feeding them to STAGE, and compare per-operator synthesized compute times using STAGE's default calibration; if the average error rises well above the reported 4.25% on comparable workload configurations, the calibration-transfer claim is falsified.

Watch this falsifier — get emailed when new claim-graph text bears on it.

If this is right

  • Designers can generate standardized execution-trace graphs for models like GPT-3, LLaMA-3, and DeepSeek-MoE under any combination of DP, TP, SP, PP, EP, and FSDP, without running the training job.
  • At 32K GPUs, a 540B-parameter dense model's graph is synthesized in about 28 minutes, and memory usage stays under 500 MB, making large-scale design-space exploration practical on a single workstation.
  • The generated graphs preserve compute-communication dependencies, so simulators can reveal overlap opportunities; the paper shows FSDP's extra communication is often hidden by compute, and that optimal parallelism shifts with HBM capacity and network bandwidth.
  • The symbolic representation extends beyond transformers: the paper shows it can express state-space models, fully-sharded tensor parallelism (a not-yet-existing strategy), and tensor-train decomposition for non-ML workloads.
  • Validation results give concrete error bounds: communication volume errors mostly under 3%, peak-memory errors 1.3-7.4% (average ~2-3% after removing CUDA init), and compute-time errors averaging 4.25% across 12 workload configurations.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • Editorial inference: The fidelity numbers come from the same H100 cluster used to calibrate the compute model, so cross-platform transferability remains untested; a natural next experiment is generating traces for a different GPU generation with default calibration and comparing to fresh measurements.
  • Editorial inference: Because graph construction is separated from numeric instantiation, the real-data plug-in feature could be extended into a workflow that scales a small set of real kernel timings to a much larger synthetic cluster—bridging measured and synthesized traces in ways the paper only sketches.
  • Editorial inference: The communication matcher's automatic mismatch detection means STAGE could double as a feasibility oracle for hypothetical parallelization strategies (such as the fully-sharded tensor parallel example) before any implementation exists, which would be a useful contribution to parallel-strategy research.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit.

Referee Report

4 major / 7 minor

Summary. STAGE is a framework that generates Chakra-format execution graphs for distributed LLM (and MoE/SSM) workloads from a symbolic tensor-graph intermediate representation. Users specify a model template and parallelization strategy; STAGE instantiates tensor-level sharding, inserts collective communications via a producer-consumer matcher, and emits numeric execution graphs for downstream simulators such as ASTRA-Sim. The paper validates generated graphs against PyTorch/Chakra traces collected on 128 H100 GPUs, reporting peak-memory agreement of roughly 97% after excluding CUDA initialization, communication-volume errors mostly below 3% in aggregate, and compute-time errors of 0.3–15.0%. It further demonstrates generation scalability to 32K GPUs in under an hour and presents several design-space-exploration case studies. The central claim is that STAGE preserves tensor-level accuracy in compute, memory, and communication while scaling to systems far beyond those used for validation.

Significance. If established, STAGE would fill a real gap: it would let system designers generate large, realistic distributed-training workloads without access to a large cluster, and it would enable principled exploration of parallelization strategies, hardware topologies, and model variants. The paper has notable strengths: the symbolic abstraction is clean and general; the communication volumes and memory footprints are largely derived structurally rather than fit to the validation traces; the framework is released open-source; and the memory validation is strong. However, the compute-time fidelity rests on operator lookup tables and a roofline model calibrated on the same H100 software stack used for validation, and the paper itself concedes the accuracy is limited to that benchmarked system. The abstract and contributions currently overstate the evidence, especially for 32K-GPU traces where no ground truth exists. These issues are fixable, but the claims must be re-scoped or additional validation must be provided.

major comments (4)
  1. [§V-C, Table VIII, Abstract] The compute-time fidelity claim is in-sample. The compute model is calibrated on the same H100 stack (NeMo 24.07, CUDA 12.5, PyTorch 2.5.0) used for validation, and the paper concedes 'we can achieve a high accuracy that limited only to the benchmarked system.' Yet the Abstract promises tensor-level accuracy for traces spanning 32K GPUs, and §VI applies the same calibrated constants to 40/80/24 GB H100 variants and different network bandwidths without re-calibration or holdout validation. Please add a holdout validation (e.g., calibrate on a subset of models/configs and predict a held-out one, ideally on a different GPU/software stack) or explicitly restrict the fidelity claim to the benchmarked stack.
  2. [Table VI, §V-C] Operator-count fidelity for the 'Others' category is poor. For example, GPT-3-5B TP=8 shows 165,207 measured vs 21,504 synthesized 'Others' ops; DeepSeek 8E shows 46,938 vs 21,952. The authors argue these are negligible in runtime, but execution-graph fidelity also depends on dependency structure, critical paths, and overlap opportunities. Missing data-movement and memory-allocation nodes can change how ASTRA-Sim schedules communication around compute. Please report structural fidelity metrics (e.g., critical-path comparison or simulation sensitivity to the missing nodes) or narrow the claim to runtime-level fidelity for dominant operators.
  3. [Table VII, §V-D] Aggregate communication-volume errors are small, but per-type volumes can be off by an order of magnitude. For Mixtral 8x7 (128 GPUs), measured vs synthesized AllReduce is 0.329 vs 16.384 MB and ReduceScatter is 13,153 vs 17,717 MB; for DeepSeek 8E, AllGather is 142 vs 3,759 MB. The footnote attributes this to micro-batch-1 expert sparsity vs STAGE's all-experts-activated default. That is not a minor issue because §VI-C explicitly targets small-batch decoding where sparse expert activation is typical. Per-type volume determines network mix and overlap. Please either validate with a batch size that activates all experts or support sparse expert activation in STAGE and report per-type errors.
  4. [§VI-B, Fig. 13, Abstract] The 32K-GPU traces are generated but not validated, yet the Abstract states they preserve tensor-level accuracy in compute, memory, and communication. Ground truth is available only for the configurations in Tables V–VIII at up to 128 GPUs. Fidelity at 32K is an extrapolation that depends on the structural communication derivation and on compute constants calibrated at 128 GPUs. Please separate generation-time scalability, which is demonstrated, from fidelity at scale, which is not, or add validation on an intermediate cluster size.
minor comments (7)
  1. [Abstract] The provided abstract says 'spanning over 128K GPUs' while the full-text abstract and §VI say 32K GPUs. Please reconcile this inconsistency.
  2. [§I, Fig. 13] The contribution says '32K GPUs less than 30 minutes,' but Fig. 13 shows 50.79 minutes for Mixtral 8x7 at 32K GPUs. Rephrase as 'under about an hour' or describe per-model times.
  3. [§VI-B, Fig. 13] §VI-B says the scalability study spans 512 to 16K GPUs, while Fig. 13's x-axis starts at 2048 and goes to 32,768. Align the text with the figure.
  4. [Table V] The 'Error Rate*' column should state the exact formula, especially how 'after removing CUDA initialization' is applied to both measured and synthesized values.
  5. [§V-C] The sentence 'we can achieve a high accuracy that limited only to the benchmarked system, however, that is sufficient enough for our validation' is ungrammatical and should be rewritten; also 'are are negligible' appears in the same section.
  6. [Abstract and general] There are several typos and broken formatting, e.g., 'publicy available', 'emgering parallel strategies', and the GitHub URL 'https://github.com/astra-sim/symbolic tensor graph' contains a space and is not a valid link.
  7. [Fig. 4] The symbolic annotation labels in Fig. 4 are very small and hard to read in a print layout; consider higher-resolution vector graphics.

Circularity Check

0 steps flagged

No significant circularity; compute-time calibration is an acknowledged scope limitation, not a derivation from the target result.

full rationale

STAGE's central derivation is structural: symbolic tensor graphs (§IV-C) are assembled from model templates and parallelization strategies, distributed via tensor-level and graph-level distributors (§IV-D), and instantiated numerically (§IV-E). Communication operators and volumes, memory lifetimes, and operator counts are consequences of tensor sharding and dependency rules, and they are validated against real Chakra traces in Tables V–VII without being fitted to those traces. The only candidate for circularity is the compute-time model in §V-C: 'we utilized a compute model with a mixture of look-up table of benchmarked operators, as well as a calibrated roofline model based on benchmarked operators.' This calibration is performed on the same H100/NVIDIA NeMo stack used for validation, so the Table VIII error rates do not demonstrate transferability to unseen GPU/software stacks. However, the paper explicitly acknowledges the limitation: 'we can achieve a high accuracy that limited only to the benchmarked system, however, that is sufficient enough for our validation.' This is an admitted scope restriction rather than a hidden reduction of the claimed result to its own inputs. The compute model composes measured operator costs and is checked on workloads (different models, batch sizes, parallelism configurations) distinct from the calibration points, so the validation is meaningful, not tautological. Self-citations to Chakra [57] and ASTRA-Sim [66] are toolchain dependencies (output format and downstream simulator), not load-bearing evidence for accuracy; no uniqueness theorem or prior-work ansatz is used to force the central design. The 32K-GPU high-fidelity claim is extrapolated without validation at that scale, but that is a correctness-risk issue, not circularity. Overall, the graph structure, memory, and communication claims are independently grounded; the compute-time claim is honestly scoped to the benchmarked system.

Axiom & Free-Parameter Ledger

1 free parameters · 5 axioms · 0 invented entities

The central claim rests on standard math/einsum semantics plus domain assumptions about communication reducibility, DAG execution, and compute-time transferability. There are no newly postulated physical entities; the 'virtual node' in the communication matcher is an abstract modeling device, not a physical claim. The main fitted parameters are the compute-calibration constants.

free parameters (1)
  • compute calibration constants (roofline peaks and operator lookup table) = not reported
    §V-C: compute time is modeled with a 'calibrated roofline model based on benchmarked operators' on the same H100 hardware used for validation; these constants determine all compute-time fidelity numbers.
axioms (5)
  • standard math Einsum/symbolic tensor shape semantics correctly represent LLM computations.
    Used throughout §IV-C; the paper assumes tensor contractions and shape inference are faithful to real PyTorch/NVIDIA operations.
  • domain assumption Every collective communication can be decomposed into Pull then Push through a virtual node (Duplicated/Partition/PartialSum).
    §IV-D2, Fig. 6. If real systems use fused or non-composable collectives, generated graphs could miss or duplicate communication operations.
  • domain assumption The execution graph is a DAG with no dynamic control flow.
    Output is 'a directed acyclic graph (DAG) with explicit operation dependencies' (§IV-A). MoE expert routing is treated statically (all experts active by default), which the authors note diverges for micro-batch=1 in Table VII footnote.
  • ad hoc to paper Kernel timings calibrated on H100/H200 transfer to other hardware scales.
    §V-C 'calibrated roofline model based on benchmarked operators'; the 32K-GPU scalability results assume these constants remain valid at scale and on other systems.
  • domain assumption Memory lifetime follows immediate garbage collection after last use.
    §V-B 'assuming garbage collection immediately thereafter'; real allocators delay reuse, which the paper acknowledges as a source of the ~1GB discrepancy.

pith-pipeline@v1.3.0-alltime-deepseek · 29826 in / 10784 out tokens · 100724 ms · 2026-08-03T22:24:03.497784+00:00 · methodology

0 comments
read the original 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

Figures

Figures reproduced from arXiv: 2511.10480 by Changhai Man, Hanjiang Wu, Huan Xu, Joongun Park, Srinivas Sridharan, Tushar Krishna.

Figure 1
Figure 1. Figure 1: Overview of STAGE model behaviors rather than enabling exploration of diverse configurations or novel parallelization strategies. • Tied to AI Model. In the arms race of AI models, there continues to be rapid evolution of LLM architec￾tures—driven by innovations such as MoEs [12], [30], attention mechanism variants [3], [15], [62], and state space models [21], aimed at improving model accuracy and training… view at source ↗
Figure 2
Figure 2. Figure 2: An Example of Execution Trace for GPU Operations. [PITH_FULL_IMAGE:figures/full_fig_p003_2.png] view at source ↗
Figure 3
Figure 3. Figure 3: STAGE Generation Flow Overview. strategies to optimize system performance and resource uti￾lization. Common training frameworks such as Megatron [56], NeMo [24], and HuggingFace Accelerate [22] frequently integrate multi-dimensional parallelization, combining data parallelism, tensor parallelism, and pipeline parallelism. Fur￾ther optimizations, such as sequence parallelism, are also increasingly adopted i… view at source ↗
Figure 4
Figure 4. Figure 4: Using Symbolic Tensor Representation to Annotate MultiHead [PITH_FULL_IMAGE:figures/full_fig_p005_4.png] view at source ↗
Figure 5
Figure 5. Figure 5: Tensor Distribution Mismatch: After applying tensor-level distribution [PITH_FULL_IMAGE:figures/full_fig_p006_5.png] view at source ↗
Figure 6
Figure 6. Figure 6: Collective Communication can be divided into two steps: Pull + Push. [PITH_FULL_IMAGE:figures/full_fig_p006_6.png] view at source ↗
Figure 7
Figure 7. Figure 7: Timing breakdown for different operators [PITH_FULL_IMAGE:figures/full_fig_p009_7.png] view at source ↗
Figure 8
Figure 8. Figure 8: Peak Memory Usage vs Runtime across configurations. [PITH_FULL_IMAGE:figures/full_fig_p010_8.png] view at source ↗
Figure 11
Figure 11. Figure 11: Memory w/ and w/o Activation Recomputation [PITH_FULL_IMAGE:figures/full_fig_p010_11.png] view at source ↗
Figure 10
Figure 10. Figure 10: Compute-Comms Overlap vs Runtime, PaLM-540B @ 64 H100 [PITH_FULL_IMAGE:figures/full_fig_p010_10.png] view at source ↗
Figure 13
Figure 13. Figure 13: STAGE Runtime Scaling with Number of GPUs TABLE X [PITH_FULL_IMAGE:figures/full_fig_p011_13.png] view at source ↗

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

79 extracted references · 38 linked inside Pith

  1. [1]

    Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model

    “Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model.” [Online]. Available: https://arxiv.org/abs/2405.04434

  2. [2]

    TensorFlow: Large-scale machine learning on heterogeneous systems,

    M. Abadi, A. Agarwal, P. Barham, E. Brevdo, Z. Chen, C. Citro, G. S. Corrado, A. Davis, J. Dean, M. Devin, S. Ghemawat, I. Goodfellow, A. Harp, G. Irving, M. Isard, Y . Jia, R. Jozefowicz, L. Kaiser, M. Kudlur, J. Levenberg, D. Man ´e, R. Monga, S. Moore, D. Murray, C. Olah, M. Schuster, J. Shlens, B. Steiner, I. Sutskever, K. Talwar, P. Tucker, V . Vanho...

  3. [3]

    Gqa: Training generalized multi-query transformer models from multi-head checkpoints,

    J. Ainslie, J. Lee-Thorp, M. de Jong, Y . Zemlyanskiy, F. Lebr ´on, and S. Sanghai, “Gqa: Training generalized multi-query transformer models from multi-head checkpoints,” 2023. [Online]. Available: https://arxiv.org/abs/2305.13245

  4. [4]

    vtrain: A simula- tion framework for evaluating cost-effective and compute-optimal large language model training,

    J. Bang, Y . Choi, M. Kim, Y . Kim, and M. Rhu, “vtrain: A simula- tion framework for evaluating cost-effective and compute-optimal large language model training,”arXiv preprint arXiv:2312.12391, 2023

  5. [5]

    Deepbench: Benchmarking json document stores,

    S. Belloni, D. Ritter, M. Schr ¨oder, and N. R ¨orup, “Deepbench: Benchmarking json document stores,” ser. DBTest ’22. New York, NY , USA: Association for Computing Machinery, 2022, p. 1–9. [Online]. Available: https://doi.org/10.1145/3531348.3532176

  6. [6]

    Language models are few-shot learners,

    T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert-V oss, G. Krueger, T. Henighan, R. Child, A. Ramesh, D. M. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amo...

  7. [7]

    Aws trainium: The journey for designing and optimization full stack ml hardware,

    N. Bshara, “Aws trainium: The journey for designing and optimization full stack ml hardware,” inProceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3, ser. ASPLOS ’24. New York, NY , USA: Association for Computing Machinery, 2024, p. 4. [Online]. Available: https://doi.org/1...

  8. [8]

    CS-3 System,

    Cerebras Systems, Inc., “CS-3 System,” https://www.cerebras.ai/system, n.d., accessed: 2025-08-01

  9. [9]

    Tvm: An automated end-to-end optimizing compiler for deep learning,

    T. Chen, T. Moreau, Z. Jiang, L. Zheng, E. Yan, M. Cowan, H. Shen, L. Wang, Y . Hu, L. Ceze, C. Guestrin, and A. Krishnamurthy, “Tvm: An automated end-to-end optimizing compiler for deep learning,” 2018. [Online]. Available: https://arxiv.org/abs/1802.04799

  10. [10]

    Llmservingsim: A simulation infrastructure for llm inference serving systems

    J. Cho, M. Kim, H. Choi, and J. Park, “Llmservingsim: A simulation infrastructure for llm inference serving systems.”

  11. [11]

    Palm: Scaling language modeling with pathways,

    A. Chowdhery, S. Narang, J. Devlin, M. Bosma, G. Mishra, A. Roberts, P. Barham, H. W. Chung, C. Sutton, S. Gehrmann, P. Schuh, K. Shi, S. Tsvyashchenko, J. Maynez, A. Rao, P. Barnes, Y . Tay, N. Shazeer, V . Prabhakaran, E. Reif, N. Du, B. Hutchinson, R. Pope, J. Bradbury, J. Austin, M. Isard, G. Gur-Ari, P. Yin, T. Duke, A. Levskaya, S. Ghemawat, S. Dev,...

  12. [12]

    Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models,

    D. Dai, C. Deng, C. Zhao, R. X. Xu, H. Gao, D. Chen, J. Li, W. Zeng, X. Yu, Y . Wu, Z. Xie, Y . K. Li, P. Huang, F. Luo, C. Ruan, Z. Sui, and W. Liang, “Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models,” 2024. [Online]. Available: https://arxiv.org/abs/2401.06066

  13. [13]

    (2025, Feb.) Deepseek v3/r1 inference system overview

    deepseek ai. (2025, Feb.) Deepseek v3/r1 inference system overview. GitHub: Open Infra Index, Day 6 of 2025 Open Source Week. Accessed: 2025-10-20. [Online]. Available: https://github.com/deepseek-ai/open- infra-index/blob/main/202502OpenSourceWeek/day 6 one more thing deepseekV3R1 inference system overview.md

  14. [14]

    Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning,

    DeepSeek-AI, D. Guo, D. Yang, H. Zhang, J. Song, R. Zhang, R. Xu, Q. Zhu, S. Ma, P. Wang, X. Bi, X. Zhang, X. Yu, Y . Wu, Z. F. Wu, Z. Gou, Z. Shao, Z. Li, Z. Gao, A. Liu, B. Xue, B. Wang, B. Wu, B. Feng, C. Lu, C. Zhao, C. Deng, C. Zhang, C. Ruan, D. Dai, D. Chen, D. Ji, E. Li, F. Lin, F. Dai, F. Luo, G. Hao, G. Chen, G. Li, H. Zhang, H. Bao, H. Xu, H. W...

  15. [15]

    Deepseek-v3 technical report,

    DeepSeek-AI, A. Liu, B. Feng, B. Xue, B. Wang, B. Wu, C. Lu, C. Zhao, C. Deng, C. Zhang, C. Ruan, D. Dai, D. Guo, D. Yang, D. Chen, D. Ji, E. Li, F. Lin, F. Dai, F. Luo, G. Hao, G. Chen, G. Li, H. Zhang, H. Bao, H. Xu, H. Wang, H. Zhang, H. Ding, H. Xin, H. Gao, H. Li, H. Qu, J. L. Cai, J. Liang, J. Guo, J. Ni, J. Li, J. Wang, J. Chen, J. Chen, J. Yuan, J...

  16. [16]

    Bert: Pre-training of deep bidirectional transformers for language understanding,

    J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova, “Bert: Pre-training of deep bidirectional transformers for language understanding,” 2019. [Online]. Available: https://arxiv.org/abs/1810.04805

  17. [17]

    Proteus: Simulating the performance of distributed dnn training,

    J. Duan, X. Li, P. Xu, X. Zhang, S. Yan, Y . Liang, and D. Lin, “Proteus: Simulating the performance of distributed dnn training,”

  18. [18]

    Available: https://arxiv.org/abs/2412.19437

    [Online]. Available: https://arxiv.org/abs/2412.19437

  19. [19]

    Zamba: A compact 7b ssm hybrid model,

    P. Glorioso, Q. Anthony, Y . Tokpanov, J. Whittington, J. Pilault, A. Ibrahim, and B. Millidge, “Zamba: A compact 7b ssm hybrid model,” 2024. [Online]. Available: https://arxiv.org/abs/2405.16712

  20. [20]

    The llama 3 herd of models,

    A. Grattafiori, A. Dubey, A. Jauhri, A. Pandey, A. Kadian, A. Al- Dahle, A. Letman, A. Mathur, A. Schelten, A. Vaughan, A. Yang, A. Fan, A. Goyal, A. Hartshorn, A. Yang, A. Mitra, A. Sravankumar, A. Korenev, A. Hinsvark, A. Rao, A. Zhang, A. Rodriguez, A. Gregerson, A. Spataru, B. Roziere, B. Biron, B. Tang, B. Chern, C. Caucheteux, C. Nayak, C. Bi, C. Ma...

  21. [21]

    Mamba: Linear-time sequence modeling with selective state spaces,

    A. Gu and T. Dao, “Mamba: Linear-time sequence modeling with selective state spaces,” 2024. [Online]. Available: https://arxiv.org/abs/ 2312.00752

  22. [22]

    Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity,

    W. Fedus, B. Zoph, and N. Shazeer, “Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity,” 2022. [Online]. Available: https://arxiv.org/abs/2101.03961

  23. [23]

    Pipedream: Fast and efficient pipeline parallel dnn training,

    A. Harlap, D. Narayanan, A. Phanishayee, V . Seshadri, N. Devanur, G. Ganger, and P. Gibbons, “Pipedream: Fast and efficient pipeline parallel dnn training,” 2018. [Online]. Available: https://arxiv.org/abs/ 1806.03377

  24. [24]

    Nemo: a toolkit for conversational ai and large language models,

    E. Harper, S. Majumdar, O. Kuchaiev, J. Li, Y . Zhang, E. Bakhturina, V . Noroozi, S. Subramanian, N. Koluguri, J. Huang, F. Jia, J. Balam, X. Yang, M. Livne, Y . Dong, S. Naren, and B. Ginsburg, “Nemo: a toolkit for conversational ai and large language models,” https://github.com/NVIDIA/NeMo, 2024, if you use this software, please cite it as above. [Onli...

  25. [25]

    Mad-max beyond single-node: Enabling large machine learning model acceleration on distributed systems,

    S. Hsia, A. Golden, B. Acun, N. Ardalani, Z. DeVito, G.-Y . Wei, D. Brooks, and C.-J. Wu, “Mad-max beyond single-node: Enabling large machine learning model acceleration on distributed systems,” in 2024 ACM/IEEE 51st Annual International Symposium on Computer Architecture (ISCA). IEEE, 2024, pp. 818–833

  26. [26]

    Gpipe: Efficient training of giant neural networks using pipeline parallelism,

    Y . Huang, Y . Cheng, A. Bapna, O. Firat, M. X. Chen, D. Chen, H. Lee, J. Ngiam, Q. V . Le, Y . Wu, and Z. Chen, “Gpipe: Efficient training of giant neural networks using pipeline parallelism,” 2019. [Online]. Available: https://arxiv.org/abs/1811.06965

  27. [27]

    Accelerate: Training and inference at scale made simple, efficient and adaptable

    S. Gugger, L. Debut, T. Wolf, P. Schmid, Z. Mueller, S. Mangrulkar, M. Sun, and B. Bossan, “Accelerate: Training and inference at scale made simple, efficient and adaptable.” https://github.com/huggingface/ accelerate, 2022

  28. [28]

    Beyond data and model parallelism for deep neural networks,

    Z. Jia, M. Zaharia, and A. Aiken, “Beyond data and model parallelism for deep neural networks,” 2018. [Online]. Available: https://arxiv.org/abs/1807.05358

  29. [29]

    Mistral 7b,

    A. Q. Jiang, A. Sablayrolles, A. Mensch, C. Bamford, D. S. Chaplot, D. de las Casas, F. Bressand, G. Lengyel, G. Lample, L. Saulnier, L. R. Lavaud, M.-A. Lachaux, P. Stock, T. L. Scao, T. Lavril, T. Wang, T. Lacroix, and W. E. Sayed, “Mistral 7b,” 2023. [Online]. Available: https://arxiv.org/abs/2310.06825

  30. [30]

    Mixtral of experts,

    A. Q. Jiang, A. Sablayrolles, A. Roux, A. Mensch, B. Savary, C. Bamford, D. S. Chaplot, D. de las Casas, E. B. Hanna, F. Bressand, G. Lengyel, G. Bour, G. Lample, L. R. Lavaud, L. Saulnier, M.-A. Lachaux, P. Stock, S. Subramanian, S. Yang, S. Antoniak, T. L. Scao, T. Gervet, T. Lavril, T. Wang, T. Lacroix, and W. E. Sayed, “Mixtral of experts,” 2024. [Onl...

  31. [31]

    Scaling laws for neural language models,

    J. Kaplan, S. McCandlish, T. Henighan, T. B. Brown, B. Chess, R. Child, S. Gray, A. Radford, J. Wu, and D. Amodei, “Scaling laws for neural language models,” 2020. [Online]. Available: https://arxiv.org/abs/2001.08361

  32. [32]

    Calculon: a methodology and tool for high-level co-design of systems and large language models,

    M. Isaev, N. McDonald, L. Dennison, and R. Vuduc, “Calculon: a methodology and tool for high-level co-design of systems and large language models,” inProceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, 2023, pp. 1–14

  33. [33]

    Gshard: Scaling giant models with conditional computation and automatic sharding,

    D. Lepikhin, H. Lee, Y . Xu, D. Chen, O. Firat, Y . Huang, M. Krikun, N. Shazeer, and Z. Chen, “Gshard: Scaling giant models with conditional computation and automatic sharding,” 2020. [Online]. Available: https://arxiv.org/abs/2006.16668

  34. [34]

    Pytorch distributed: Experiences on accelerating data parallel training,

    S. Li, Y . Zhao, R. Varma, O. Salpekar, P. Noordhuis, T. Li, A. Paszke, J. Smith, B. Vaughan, P. Damania, and S. Chintala, “Pytorch distributed: Experiences on accelerating data parallel training,” 2020. [Online]. Available: https://arxiv.org/abs/2006.15704

  35. [35]

    Mystique: Enabling accurate and scalable generation of production ai benchmarks,

    M. Liang, W. Fu, L. Feng, Z. Lin, P. Panakanti, S. Zheng, S. Sridharan, and C. Delimitrou, “Mystique: Enabling accurate and scalable generation of production ai benchmarks,” 2023. [Online]. Available: https://arxiv.org/abs/2301.04122

  36. [36]

    Lumos: Efficient performance modeling and estimation for large-scale llm training,

    M. Liang, H. T. Kassa, W. Fu, B. Coutinho, L. Feng, and C. Delimitrou, “Lumos: Efficient performance modeling and estimation for large-scale llm training,” 2025. [Online]. Available: https://arxiv.org/abs/2504.09307

  37. [37]

    Reducing activation recomputation in large transformer models,

    V . Korthikanti, J. Casper, S. Lym, L. McAfee, M. Andersch, M. Shoeybi, and B. Catanzaro, “Reducing activation recomputation in large transformer models,” 2022. [Online]. Available: https: //arxiv.org/abs/2205.05198

  38. [38]

    (2025, Apr.) The llama 4 herd: The beginning of a new era of natively multimodal ai innovation

    Meta AI. (2025, Apr.) The llama 4 herd: The beginning of a new era of natively multimodal ai innovation. Accessed: 2025-04-22. [Online]. Available: https://ai.meta.com/blog/llama-4-multimodal-intelligence/

  39. [39]

    Chakra working group,

    MLCommons, “Chakra working group,” https://mlcommons.org/ working-groups/research/chakra/, 2023

  40. [40]

    Chakra schema,

    MLCommons, “Chakra schema,” https://github.com/mlcommons/ chakra/wiki/Chakra-Schema, 2024

  41. [41]

    Mini- batch coresets for memory-efficient training of large language models,

    D. Nguyen, W. Yang, R. Anand, Y . Yang, and B. Mirzasoleiman, “Mini- batch coresets for memory-efficient training of large language models,” inarXiv:2407.19580 [cs.LG], 2024

  42. [42]

    Jamba: A hybrid transformer-mamba language model,

    O. Lieber, B. Lenz, H. Bata, G. Cohen, J. Osin, I. Dalmedigos, E. Safahi, S. Meirom, Y . Belinkov, S. Shalev-Shwartz, O. Abend, R. Alon, T. Asida, A. Bergman, R. Glozman, M. Gokhman, A. Manevich, N. Ratner, N. Rozen, E. Shwartz, M. Zusman, and Y . Shoham, “Jamba: A hybrid transformer-mamba language model,”

  43. [43]

    Available: https://arxiv.org/abs/2403.19887

    [Online]. Available: https://arxiv.org/abs/2403.19887

  44. [44]

    Nvidia nemo - open-source toolkit for conversational ai,

    NVIDIA, “Nvidia nemo - open-source toolkit for conversational ai,” 2024, accessed: 2024-11-23. [Online]. Available: https://www.nvidia. com/en-us/ai-data-science/products/nemo/

  45. [45]

    Cuda profiling tools interface (cupti),

    NVIDIA Corporation, “Cuda profiling tools interface (cupti),” https:// developer.nvidia.com/cupti, 2024, accessed: 2024-11-21

  46. [46]

    NVIDIA HGX Platform,

    NVIDIA Corporation, “NVIDIA HGX Platform,” https://www.nvidia. com/en-us/data-center/hgx/, n.d., accessed: 2025-08-01

  47. [47]

    Tensor-train decomposition,

    I. V . Oseledets, “Tensor-train decomposition,”SIAM Journal on Scien- tific Computing, vol. 33, no. 5, pp. 2295–2317, 2011

  48. [48]

    Dnnfusion: accelerating deep neural networks execution with advanced operator fusion,

    W. Niu, J. Guan, Y . Wang, G. Agrawal, and B. Ren, “Dnnfusion: accelerating deep neural networks execution with advanced operator fusion,” inProceedings of the 42nd ACM SIGPLAN International Conference on Programming Language Design and Implementation, ser. PLDI 2021. New York, NY , USA: Association for Computing Machinery, 2021, p. 883–898. [Online]. Ava...

  49. [49]

    Nvidia cupti - cuda profiling tools interface,

    NVIDIA, “Nvidia cupti - cuda profiling tools interface,” 2024, accessed: 2024-11-23. [Online]. Available: https://developer.nvidia.com/cupti

  50. [50]

    Kineto: Performance profiling library for pytorch,

    PyTorch Team, “Kineto: Performance profiling library for pytorch,” https://github.com/pytorch/kineto, 2025, accessed: 2025-07-31

  51. [51]

    Zero: Memory optimizations toward training trillion parameter models,

    S. Rajbhandari, J. Rasley, O. Ruwase, and Y . He, “Zero: Memory optimizations toward training trillion parameter models,” inSC20: International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 2020, pp. 1–16

  52. [52]

    Themis: A Network Bandwidth-Aware Collective Scheduling Policy for Distributed Training of DL Models,

    S. Rashidi, W. Won, S. Srinivasan, S. Sridharan, and T. Krishna, “Themis: A Network Bandwidth-Aware Collective Scheduling Policy for Distributed Training of DL Models,” inProceedings of the 49th Annual International Symposium on Computer Architecture (ISCA ’22), 2022, p. 581–596

  53. [53]

    Mlperf inference benchmark,

    V . J. Reddi, C. Cheng, D. Kanter, P. Mattson, G. Schmuelling, C.-J. Wu, B. Anderson, M. Breughe, M. Charlebois, W. Chou, R. Chukka, C. Coleman, S. Davis, P. Deng, G. Diamos, J. Duke, D. Fick, J. S. Gardner, I. Hubara, S. Idgunji, T. B. Jablin, J. Jiao, T. S. John, P. Kanwar, D. Lee, J. Liao, A. Lokhmotov, F. Massa, P. Meng, P. Micikevicius, C. Osborne, G...

  54. [54]

    Pytorch: An imperative style, high-performance deep learning library,

    A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. K ¨opf, E. Yang, Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala, “Pytorch: An imperative style, high-performance deep learning library,” 2019. [Online]. Available: https://arxiv.org...

  55. [55]

    Pytorch profiler recipe,

    PyTorch Contributors, “Pytorch profiler recipe,” https://pytorch.org/ tutorials/recipes/recipes/profiler recipe.html, 2024, accessed: 2024-11- 21

  56. [56]

    Megatron-lm: Training multi-billion parameter language models using model parallelism,

    M. Shoeybi, M. Patwary, R. Puri, P. LeGresley, J. Casper, and B. Catanzaro, “Megatron-lm: Training multi-billion parameter language models using model parallelism,” 2020. [Online]. Available: https://arxiv.org/abs/1909.08053

  57. [57]

    Chakra: Advancing performance benchmarking and co-design using standardized execution traces,

    S. Sridharan, T. Heo, L. Feng, Z. Wang, M. Bergeron, W. Fu, S. Zheng, B. Coutinho, S. Rashidi, C. Manet al., “Chakra: Advancing performance benchmarking and co-design using standardized execution traces,”arXiv preprint arXiv:2305.14516, 2023

  58. [58]

    Kineto: A cpu+gpu profiling library for pytorch,

    P. Team, “Kineto: A cpu+gpu profiling library for pytorch,” https: //github.com/pytorch/kineto, 2023, accessed: 2025-04-14

  59. [59]

    Llama: Open and efficient foundation language models,

    H. Touvron, T. Lavril, G. Izacard, X. Martinet, M.-A. Lachaux, T. Lacroix, B. Rozi `ere, N. Goyal, E. Hambro, F. Azhar, A. Rodriguez, A. Joulin, E. Grave, and G. Lample, “Llama: Open and efficient foundation language models,” 2023. [Online]. Available: https: //arxiv.org/abs/2302.13971

  60. [60]

    Torch.fx: Practical program capture and transformation for deep learning in python,

    J. K. Reed, Z. DeVito, H. He, A. Ussery, and J. Ansel, “Torch.fx: Practical program capture and transformation for deep learning in python,” 2022. [Online]. Available: https://arxiv.org/abs/2112.08429

  61. [61]

    Param: A trace abstraction for ml workloads,

    F. Research, “Param: A trace abstraction for ml workloads,” https:// github.com/facebookresearch/param, 2023, accessed: 2025-04-14

  62. [63]

    Attention is all you need,

    A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin, “Attention is all you need,” 2023. [Online]. Available: https://arxiv.org/abs/1706.03762

  63. [64]

    Rail- only: A Low-Cost High-Performance Network for Training LLMs with Trillion Parameters ,

    W. Wang, M. Ghobadi, K. Shakeri, Y . Zhang, and N. Hasani, “ Rail- only: A Low-Cost High-Performance Network for Training LLMs with Trillion Parameters ,” inProceedings of the 2024 IEEE Symposium on High-Performance Interconnects (HOTI), 2024

  64. [65]

    SimAI: Unifying architecture design and performance tuning for Large-Scale large language model training with scalability and precision,

    X. Wang, Q. Li, Y . Xu, G. Lu, D. Li, L. Chen, H. Zhou, L. Zheng, S. Zhang, Y . Zhu, Y . Liu, P. Zhang, K. Qian, K. He, J. Gao, E. Zhai, D. Cai, and B. Fu, “SimAI: Unifying architecture design and performance tuning for Large-Scale large language model training with scalability and precision,” in22nd USENIX Symposium on Networked Systems Design and Implem...

  65. [66]

    Unity: Accelerating DNN training through joint optimization of algebraic transformations and parallelization,

    C. Unger, Z. Jia, W. Wu, S. Lin, M. Baines, C. E. Q. Narvaez, V . Ramakrishnaiah, N. Prajapati, P. McCormick, J. Mohd-Yusof, X. Luo, D. Mudigere, J. Park, M. Smelyanskiy, and A. Aiken, “Unity: Accelerating DNN training through joint optimization of algebraic transformations and parallelization,” in16th USENIX Symposium on Operating Systems Design and Impl...

  66. [67]

    Tensor comprehensions: Framework-agnostic high- performance machine learning abstractions,

    N. Vasilache, O. Zinenko, T. Theodoridis, P. Goyal, Z. DeVito, W. S. Moses, S. Verdoolaege, A. Adams, and A. Cohen, “Tensor comprehensions: Framework-agnostic high- performance machine learning abstractions,” 2018. [Online]. Available: https://arxiv.org/abs/1802.04730

  67. [68]

    Optimizing deep learning inference via global analysis and tensor expressions,

    C. Xia, J. Zhao, Q. Sun, Z. Wang, Y . Wen, T. Yu, X. Feng, and H. Cui, “Optimizing deep learning inference via global analysis and tensor expressions,” inProceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 1, ser. ASPLOS ’24. New York, NY , USA: Association for Computing Mac...

  68. [69]

    Apollo: Automatic partition-based operator fusion through layer by layer optimization,

    J. Zhao, X. Gao, R. Xia, Z. Zhang, D. Chen, L. Chen, R. Zhang, Z. Geng, B. Cheng, and X. Jin, “Apollo: Automatic partition-based operator fusion through layer by layer optimization,” inProceedings of Machine Learning and Systems, D. Marculescu, Y . Chi, and C. Wu, Eds., vol. 4, 2022, pp. 1–

  69. [70]

    Pytorch fsdp: Experiences on scaling fully sharded data parallel,

    Y . Zhao, A. Gu, R. Varma, L. Luo, C.-C. Huang, M. Xu, L. Wright, H. Shojanazeri, M. Ott, S. Shleifer, A. Desmaison, C. Balioglu, P. Damania, B. Nguyen, G. Chauhan, Y . Hao, A. Mathews, and S. Li, “Pytorch fsdp: Experiences on scaling fully sharded data parallel,”

  70. [71]

    Mist: Efficient distributed training of large language models via memory-parallelism co-optimization,

    Z. Zhu, C. Giannoula, M. Andoorveedu, Q. Su, K. Mangalam, B. Zheng, and G. Pekhimenko, “Mist: Efficient distributed training of large language models via memory-parallelism co-optimization,” inProceedings of the Twentieth European Conference on Computer Systems, ser. EuroSys ’25. ACM, Mar. 2025, p. 1298–1316. [Online]. Available: http://dx.doi.org/10.1145...

  71. [72]

    Astra-sim2. 0: Modeling hierarchical networks and disaggregated systems for large-model training at scale,

    W. Won, T. Heo, S. Rashidi, S. Sridharan, S. Srinivasan, and T. Kr- ishna, “Astra-sim2. 0: Modeling hierarchical networks and disaggregated systems for large-model training at scale,” in2023 IEEE International Symposium on Performance Analysis of Systems and Software (ISPASS). IEEE, 2023, pp. 283–294

  72. [73]

    Astra-sim2.0: Modeling hierarchical networks and disaggregated systems for large-model training at scale,

    W. Won, T. Heo, S. Rashidi, S. Sridharan, S. Srinivasan, and T. Krishna, “Astra-sim2.0: Modeling hierarchical networks and disaggregated systems for large-model training at scale,” 2023. [Online]. Available: https://arxiv.org/abs/2303.14006

  73. [76]

    Available: https://proceedings.mlsys.org/paper files/paper/ 2022/file/e175e8a86d28d935be4f43719651f86d-Paper.pdf

    [Online]. Available: https://proceedings.mlsys.org/paper files/paper/ 2022/file/e175e8a86d28d935be4f43719651f86d-Paper.pdf

  74. [78]

    Available: https://arxiv.org/abs/2304.11277

    [Online]. Available: https://arxiv.org/abs/2304.11277

  75. [80]

    Resiliency at scale: Managing Google’s TPUv4 machine learning supercomputer,

    Y . Zu, A. Ghaffarkhah, H.-V . Dang, B. Towles, S. Hand, S. Huda, A. Bello, A. Kolbasov, A. Rezaei, D. Du, S. Lacy, H. Wang, A. Wisner, C. Lewis, and H. Bahini, “Resiliency at scale: Managing Google’s TPUv4 machine learning supercomputer,” in21st USENIX Symposium on Networked Systems Design and Implementation (NSDI 24). Santa Clara, CA: USENIX Association...

  76. [2020]

    Available: https://arxiv.org/abs/2005.14165

    [Online]. Available: https://arxiv.org/abs/2005.14165

  77. [2023]

    Available: https://arxiv.org/abs/2306.02267

    [Online]. Available: https://arxiv.org/abs/2306.02267

  78. [2024]

    Available: https://arxiv.org/abs/2407.21783

    [Online]. Available: https://arxiv.org/abs/2407.21783

  79. [2025]

    Available: https://arxiv.org/abs/2501.12948

    [Online]. Available: https://arxiv.org/abs/2501.12948