arXiv CSopen access

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

A Few GPUs, A Whole Lotta Scale: Faithful LLM Training Emulation with PrismLLM Shaoke Xi1,∗ , ChonLam Lao1,2,∗ , Boyi Jia1,3,∗ , Jiaqi Gao1,∗ , Zhipeng Zhang1 , Jiamin Cao1 , Brian Sutioso2 , Erci Xu3 , Minlan Yu2 , Kui Ren4 , Yong Li1 , Zhengping Qian1 , Ennan Zhai1 , Jingren Zhou1 1 Alibaba Group, 2 Harvard University, 3 Shanghai Jiao Tong University, 4 Zhejiang University

arXiv:2605.15617v1 [cs.DC] 15 May 2026

ABSTRACT

of developing, debugging, and optimizing the systems that enable these workloads. In practice, LLM training follows a DevOps workflow in which engineers continuously analyze training telemetry, implement optimizations, and deploy updates during ongoing runs. This workflow critically depends on the ability to reproduce and analyze field behavior under realistic conditions. Unfortunately, as the cluster scale grows, the cost of system experimentation increases proportionally. Reproducing a performance issue, validating a design change, or conducting a what-if analysis often requires access to the same production-scale cluster that runs the training job [12, 26]. Allocating thousands of GPUs solely for debugging or experimentation is both expensive and operationally disruptive, making rapid system iteration increasingly difficult. Two existing solutions partially alleviate the issue, but neither suffices. One option is simulation [1, 9, 10, 13, 21, 27, 28, 33], which lowers hardware cost by replacing real execution with profiled traces and analytical models. But simulation accuracy depends on the fidelity of those models, and maintaining them is increasingly difficult as training stacks evolve rapidly across frameworks, compilers, kernels, network, and hardware generations [4, 5, 11, 18, 23, 30]. The other option is to run real stack at smaller scale. However, downscaling changes exactly the behaviors engineers care about: world size, parallelization strategy, batch size, and scheduling choices all reshape communication structure, memory layout, and bottlenecks [20, 24]. Hence, conclusions drawn from downscaled experiments often fail to transfer to production. In short, we lack a way to study large-scale training behavior faithfully without paying excessive cost. This paper presents PrismLLM, a system that enables faithful emulation of large-scale LLM training using limited hardware resources. The key insight behind PrismLLM is that large-scale interaction structure does not require all ranks to execute simultaneously on physical GPUs. Instead, scale can be virtualized by capturing the execution structure of a training job and selectively replaying it while running only a subset of ranks on real hardware. By decoupling the logical training scale from physical GPU allocation, PrismLLM allows engineers to observe large-scale behavior while using only a small fraction of the original hardware.

Large language model (LLM) training today runs on clusters spanning thousands of GPUs. While this scale enables rapid model advances, developing, debugging, and performancetuning, 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 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.

1

INTRODUCTION

Training large language models (LLMs) has been scaling at an unprecedented speed. Production systems today routinely involve thousands to tens of thousands of GPUs [6, 8, 14, 17, 29], and emerging frontier models are expected to push this scale even further. While significant research has focused on improving training performance at such levels [4, 5, 22– 24, 30, 32], much less attention has been paid for the cost ∗ These authors contributed equally to this work.

1

PrismLLM achieves this through a two-phase design. First, PrismLLM constructs a high-fidelity execution graph that captures computation, communication, and dependency relationships across all ranks of a training job. This graph is generated using a context-switching execution mechanism that multiplexes many logical ranks onto a small set of GPUs. Second, PrismLLM performs hybrid emulation in which selected ranks execute the real training program on physical GPUs while the remaining ranks are replayed as virtual participants that preserve the communication and synchronization behavior of the original deployment. This design allows engineers to run unmodified training code while observing realistic, large-scale interactions. We evaluate PrismLLM on large-scale LLM training workloads with up to thousands of GPUs. Our results show that PrismLLM accurately reproduces both performance and memory behavior of production-scale training runs while requiring only a small fraction of the original hardware. Across a diverse set of models and parallelization strategies, PrismLLM predicts iteration time with an average error of 0.58% and reproduces GPU memory usage with negligible error. Moreover, PrismLLM enables practical experimentation and debugging workflows that would otherwise require dedicating an entire production cluster. We commit to open source PrismLLM and our traces to expedite future research efforts and industry adoption.

2

memory pressure, and scale-dependent instability, surface only at larger scale. • Production rollout. New modifications enter the production training job only at permitted points such as checkpoint boundaries or restarts after fail-stop events. Engineers then monitor the live system for throughput, efficiency, and stability improvements. If regressions appear, they roll back to a stable checkpoint and return to smaller-scale environments for debugging before retrying. Reproducing scale-dependent failures sometimes requires temporarily dedicating the entire production cluster to debugging [26].

3

MOTIVATION

3.1

The Dilemma of DevOps in LLM Training

Under the DevOps paradigm, engineers are expected to devise new techniques and closely monitor the effectiveness of newly-applied changes in the field. While this is critical to the success of pre-training, the developers often face a fundamental constraint: GPU scarcity. Under the fierce competition of releasing new frontier models every 3-6 months [7], the majority of GPUs are occupied by production workloads, leaving only a small fraction (hundreds or even tens) for individual framework development, debugging, and validation. As a result, the engineers can implement and locally validate changes quickly, but actual confirmation often arrives only after delayed large-scale deployment. Moreover, the access to large-scale resources is a musthave for two reasons. First, in training, performance is determined by interactions across multiple system layers, including GPU micro-architecture, communication, and the caching and memory hierarchy. Any of these components can become a bottleneck, so they must be carefully designed and optimized together. For example, engineers pay much attention to overlapping communication and computation and balancing the compute–memory ratio to maximize the utilization of expensive GPUs [2, 31]. While such optimizations are important, they often require fine-grained, lowlevel development, and code complexity grows rapidly. It is challenging for engineers to develop and maintain systems without a full picture at scale. The stack also evolves rapidly. On the software side, ML frameworks, compiler backends, and communication libraries are continuously optimized, and kernels are frequently redesigned and fused [4, 5, 23, 30]. Execution behavior becomes a constantly moving target. On the hardware side, platforms also evolve rapidly. New GPU architectures emerge roughly every two years—NVIDIA introduced A100 in 2020, H100 in 2022, and Blackwell B100 in 2024. Each generation introduces changes in compute capability, memory systems, and interconnects, often leading to significantly different performance characteristics and design spaces for engineers.

LLM TRAINING DEVOPS 101

LLM training engineers usually operate in a DevOps model where they are responsible for deploying and monitoring training jobs. In practice, the development process follows a recurring loop: • Analyze and set targets. Engineers study telemetry and failures to diagnose inefficiencies from prior or ongoing training runs. From this analysis, they formulate optimization targets, such as lower step time, better communication efficiency, reduced memory overhead, and stronger fault tolerance. These targets evolve continuously as new bottlenecks gradually unfold. This stage is particularly challenging because of the tight coupling between model execution and hardware topology: even small design changes can shift communication patterns, synchronization behavior, and memory pressure at scale. • Implement. A single optimization often spans multiple layers of the stack, from Python orchestration and distributed runtime logic to communication libraries, CUDA kernels, and observability tools. The main difficulty is coordinating changes across tightly coupled interfaces while preserving correctness, performance, and stability. • Small-scale verification. Engineers validate correctness, rule out obvious regressions, and collect preliminary performance signals on a limited number of GPUs. This stage provides only partial confidence because many problems, such as synchronization anomalies, collective contention, 2

3.2

Simulating Large-scale Clusters?

Moreover, even when carefully constructed, downscaled experiments may fail to accurately reflect large-scale performance and program behavior. Training performance is highly scale-sensitive: as scale changes, system bottlenecks can shift (e.g., from computation-bound to memory-bound). As a result, conclusions drawn from downscaled experiments may not generalize to production deployments and can even mislead optimization decisions. A really tricky example is CPU activation offloading, a common technique to reduce active memory footprint. We need to decide which activation tensors, when to offload, and when to onload to keep the memory within limit, while minimizing the conflict with the ongoing computation and communication. Since the downscaled experiment alters the memory layout, local computation workload, and pipeline schedule, the offloading behavior and benefit are no longer align with the online environment.

Engineers have long sought methods to evaluate the performance of ML training workloads without requiring largescale GPU deployments. Prior work has explored simulationbased systems [1, 9, 10, 13, 19, 21, 27, 28, 33]. These systems approximate large-scale training behavior by constructing and replaying intermediate representations using techniques such as execution trace collection, operator-level profiling, and analytical modeling, enabling performance estimation without full-scale execution. However, simulation accuracy depends on the fidelity of the underlying models: more detailed models better capture system behavior but require substantially greater engineering effort. As ML systems become increasingly complex and evolve rapidly, maintaining accurate simulations demands deep system expertise and continuous re-engineering. SimAI [27] re-implements a mocked version of the ML framework, doubling the developer’s engineering effort. Phantora [19], though it simulates at a lower level, still requires manually modeling each CUDA operation and network behavior. Also, many black-boxed behaviors, such as GPU kernel scheduling and GPU overheat throttling, are too complex to be modeled accurately. Moreover, assumptions embedded in simulation models can quickly become invalid, requiring repeated profiling and substantial re-implementation [19], making simulation increasingly costly and less practical. 3.3

3.4

Selective GPU Execution: A Promising Idea

Field development workflows suggest a different intuition. When debugging or analyzing large-scale training runtime, engineers rarely need to inspect the behavior of all ranks simultaneously. Instead, they typically focus on a small subset of ranks of interest (e.g., 8 GPUs), while the remaining ranks simply participate in the distributed execution. For example, engineers often examine the end-to-end step time by observing when the last rank completes an iteration, or inspect the execution path of a specific rank to understand program logic or diagnose failures. Similarly, debugging issues such as out-of-memory (OOM) errors or communication stalls typically requires detailed inspection of only a few ranks rather than the entire deployment. This observation raises a potential opportunity: can we execute only the ranks of interest physically while maintaining the rest of the system logically? Such an approach could preserve deployment-scale interaction structures while significantly reducing hardware requirements. Ideally, we can build a hybrid emulator that decouples logical scale from physical hardware. By collecting per-rank execution graph offline, we would only need to calibrate a few (say 8) GPUs to achieve a faithful and accurate representation (e.g., timing and dependencies). Then, engineers can specify a subset of ranks of interest to execute the real training program on physical GPUs. These selected ranks run on real hardware (i.e., calibrated GPUs), while the remaining logical ranks are virtualized to replay the execution graph to preserve the interaction structure of the full deployment.

Downscaling the Testing Environment?

Another option is downscaling the execution, which means we still run the full training software stack, but on a reduced deployment to approximate production-scale behaviors. This is typically achieved by shrinking model parameters, batch sizes, or degrees of parallelism (e.g., PP, TP, and EP) while preserving the original execution paths as much as possible. However, such approaches often require heavy engineering investment. Note that constructing a representative and faithful downscaled environment is rarely straightforward. Suppose we want to observe the behavior of pipeline parallelism between two middle stages to improve communication–computation overlap. Without access to the same number of GPUs as in production, we must reduce other parallelism dimensions (DP, EP, and TP) while preserving a pipeline configuration. Many related parameters must also be reconfigured to fit the new setup, such as the virtual pipeline parallelism (VPP) degree and the global batch size. This process fundamentally alters the training configuration, and thus such modifications are unlikely to be useful for actual deployment. A typical example is VPP’s pipeline scheduling. In the small scale, we have to reduce the batch size to avoid memory overflow. We need to specifically handle the case when the batch size is smaller than the VPP since there is no enough data to drive the pipeline. However, such case is impossible in the online setting.

3.5

Envision: A Trace-driven Emulator

Based on our in-production experience, we envision an emulation system that achieves three goals: • Low resource footprint. The system should emulate the large-scale workload with only a few GPUs. • High fidelity. The system should faithfully reproduce the behavior of the ranks of interest as large-scale runs. 3

Stage 2 Timing Filling and Calibration

Stage 1 Graph Collection

②

Bare Bare Graph Graph 0 1

Bare Bare Graph Graph 2 3

Execution Execution Graph Graph

Data Data Data

Replay Replay Replay Replay Replay Replay Graph 100 Replay Replay ReplayvRank vRank vRank 103 103 103 vRank vRank vRank 103 103 103 NCCL NCCL NCCL vRank vRank 103 103 vRank 100

Data Rank 0

Rank 1

Rank 2

Rank 3

All GPUs

Rank NC 1 CL

Rank NC 2 CL

Rank 3 NC CL

Assistant GPU Machine

④

Sandbox GPU Machine

⑥

①

Rank Program Rank 1024 Rank 20 2 Coordinator

⑤ ·

Rank NC 0 CL

Graph Generation Subsystem (§5)

Emulation

③ Slice

⑦ Graph Graph 0 1

Graph Graph 2 3

Calibration

⑧

Execution Execution Graph Graph

Data Data Data

Replay Replay Replay Replay Replay Replay Graph 100 Replay Replay ReplayvRank vRank vRank 103 103 103 vRank vRank vRank 103 103 103 NCCL NCCL NCCL vRank vRank 103 103 vRank 100

⑨

Rank 0 NC CL

Rank 1 NC CL

Rank NC 2 CL

Rank 3 NC CL

Assistant GPU Machine Sandbox GPU Machine

⑩

Metrics

Hybrid Emulation Subsystem (§6)

Figure 1. System overview

• Code reuse. The emulation system should directly reuse the current LLM training code base to avoid unnecessary development or maintenance overhead.

graph, which is later used to replay execution without requiring large-scale hardware. In the emulation phase, PrismLLM constructs an emulation environment that allows a subset of ranks to run as if they were in a large-scale deployment. Each phase corresponds to a subsystem of PrismLLM. Phase 1: Graph preparation (§5). The first preparation subsystem constructs a high-fidelity execution representation (i.e., an execution graph) for all ranks in the program. 1 first executes the origPrismLLM’s centralized coordinator ○ inal program on a small number of GPUs by scheduling only a subset of ranks at a time and coordinating their progress. It multiplexes execution across ranks and records an execution graph for each rank. This process produces a structural graph (bare graph) that encodes operations and dependencies 2 without accurate timing ○. To obtain accurate timing under limited GPU resources, PrismLLM collects execution graph timings locally within 3 using available GPUs, and then merges and each slice ○ calibrates them to construct a globally consistent execution timeline. The system executes slices one by one. In each slice, the available GPUs are partitioned into two roles: sandbox 4 and assistant GPUs ○. 5 A subset of ranks runs as real GPUs ○ ranks on sandbox GPUs, while the remaining are replayed as virtual ones on assistant GPUs to serve as communication counterparts and preserve correct timing (e.g., responding to collective communication). PrismLLM iterates over these subsets (slices) in a round-robin manner (e.g., 0-7, 8-15, . . . ), ensuring that each rank is executed once as a real rank. 6 PrismLLM After collecting per-slice execution timings ○, performs cross-slice calibration to reconstruct the complete high-fidelity execution graph that captures globally consis7 tent large-scale behavior ○. Phase 2: Hybrid emulation (§6). After graph preparation, PrismLLM performs emulation in the second subsystem 8 In this phase, only the specified ranks of interest are exe○. cuted using the original program on sandbox GPUs, serving as the primary observation points. The remaining ranks are instantiated as virtual ranks and replayed on assistant GPUs using the constructed execution graph; these virtual ranks 9 Unlike perform real communication with sandbox ranks ○. the previous, this phase runs a single emulation instance

A natural approach is record-and-replay where we record execution trace of a program and replay it in the test environment. But, there are two challenges. Challenge 1: You need scale to see scale. Our goal is to emulate large-scale behavior using only a small number of GPUs. However, doing so requires first capturing largescale behavior for later replay, while capturing such behavior itself requires running at large scale. This creates a chickenand-egg problem: we need large-scale execution to obtain a faithful representation, yet our goal is precisely to avoid running at scale. Furthermore, even if we can somehow obtain a faithful representation of large-scale behavior, replaying a large number of logical ranks on a small number of GPUs remains challenging. Challenge 2: Preserving large-scale behavior at small scale. In large-scale settings, all ranks execute concurrently and make progress together. When we emulate on a small number of GPUs, this concurrency is lost, and interaction patterns—such as collective communication, synchronization, and computation–communication overlap—can deviate from those in full-scale execution. A core challenge is therefore to preserve these interaction patterns and their timing relationships faithfully, even when physical concurrency is drastically reduced. In particular, timing distortions introduced by resource multiplexing must be understood and corrected to maintain deployment-scale fidelity.

4

PRISMLLM DESIGN OVERVIEW

Our goal is to build a low-resource-footprint and highfidelity emulation system that allows engineers to run their existing code while executing only the ranks of interest in large-scale training, enabling efficient performance evaluation and faithful reproduction of program behavior. We hence develop PrismLLM which consists of two phases, preparation and emulation. In the preparation phase, the system executes the original large-scale program on a small number of GPUs to generate a global, high-fidelity execution 4

Rank 0

BWD#0 (6.4ms)

send#16 (7.9ms)

recv#17 (3.0ms)

FWD#9 (2.7ms)

send#18 (2.6ms)

recv#19 (2.8ms)

Context Switching Based

BWD#1 (5.6ms)

GPU Slots

…

R0

…

Rank 1

send#23 (2.6ms)

BWD#1 (6.4ms)

recv#24 (1.8ms)

send#25 (0.1ms)

FWD#7 (2.6ms)

recv#26 (0.1ms)

send#27 (2.6ms)

R1

Comm. Tensors

R2

BWD#2 (5.7ms)

R3

R8

R10

GPU Slots

R11

R0 R1 R2 R3 R4 R5

R4 R5 R6 R7 R8 R9 R10 R11

Figure 2. Snippet from PrismTrace capturing dependencies of 1F1B

Comm. Tensors

R6

R7

R9

Tensor Rules GPU Slots

where all virtual ranks are replayed using the complete execution graph with accurate timing. This enables PrismLLM to accurately measure metrics of interest (e.g., end-to-end iteration time and GPU memory usage over time) for sandbox 10 ranks without requiring full-scale hardware ○. Both phases operate under minimal GPU resources, requiring as few as two GPU machines: one hosts the sandbox GPUs running the original program, while the other hosts the assistant GPUs for replay. The number of assistant GPUs must match that of the sandbox GPUs to ensure sufficient network bandwidth.

R0

R1

R2

R3

Tensors Rules

R0

R1

R2

R3

Tensors Rules

R4 R5 R6 R7 R8 R9 R10 R11

R8 R9 R10 R11

Figure 3. Context switch to generate graph.

To enable replay during emulation, we design a replay-oriented graph representation (§5.1) that captures only essential information while faithfully preserving large-scale behavior at low cost. We construct this graph using a small number of GPUs (§5.2) and ensure accurate timing (§5.3).

where one operation must complete before another begins (e.g., communication followed by computation within a rank), and (2) synchronization, where all participating nodes must reach the operation before any can proceed (e.g., collectives or matched send–receive pairs). Finally, PrismTrace records only GPU-side communication timing, ignoring CPU-side events. Because actual communication timing is determined by GPU execution—not by when the CPU issues the call—CPU-side timestamps are misleading for replay. By focusing exclusively on execution structure and communication boundaries, PrismLLM captures only what is necessary for faithful replay, significantly reducing tracing overhead.

5.1 PrismTrace Execution Graph PrismLLM constructs an execution graph that captures both communication and computation across all ranks, providing a complete view for replay. For each rank, the graph must answer three questions: (1) what operations are executed (computation and communication spans), (2) in what order they execute (intra- and inter-rank dependencies), and (3) how long they take (duration of each span). A common approach is to use the PyTorch Profiler [18] with tools such as Chakra [25] and Meta HTA [15]. However, PyTorch Profiler incurs significant overhead—we observe up to 20.04% iteration time increase in a 128-GPU Qwen3 training job [29]—because it is not designed for replay. It captures fine-grained, operator-level events that are unnecessary for our purpose. Since our goal is to determine when communication should be issued during replay (not to re-execute computation), only communication timing and large-scale dependencies need to be preserved; operator-level details are redundant. Motivated by this insight, we design PrismTrace, a new tracing format that captures only the minimal information required for replay. PrismTrace operates at a coarser scheduling granularity (e.g., per microbatch), where each scheduling unit defines the minimal dependency boundary across downstream communication. Figure 2 shows a snippet of the graph. It consists of nodes and edges: each node represents a computation span or a communication event, identified by a unique id and a duration; each edge encodes a dependency between two nodes (id1 , id2 ). We distinguish two dependency types: (1) directional,

5.2 Stage 1: Collecting PrismTrace Collecting PrismTrace for all ranks requires running the full training job at scale, but this conflicts with the requirement of using only a small number of GPUs. To address this, PrismLLM introduces a context-switching execution system that decouples logical execution from physical scale, enabling graph collection on a small number of GPUs. At a high level, PrismLLM maintains a pool of logical ranks, each with its own execution state. At any moment, only 𝑁 ranks are active on the 𝑁 available GPUs, while the remaining ranks are paused in CPU memory (or spilled to disk). The PrismLLM coordinator schedules subsets of ranks onto the GPUs and runs them until they block on a communication point. When a rank blocks, its execution state and communication context are saved, and the system switches to other runnable ranks. By iteratively advancing all ranks in this way, PrismLLM captures the full execution graph without requiring large-scale hardware. Because each rank executes the original program with real tensor values, this approach guarantees that the program follows the correct code paths and handles all value-dependent logic faithfully. End-to-end workflow. PrismLLM runs unmodified training programs transparently and manages all ranks through a centralized coordinator. The input can be any large-scale training job (e.g., 256 GPUs with PP=4, TP=4, DP=16), as the parallelism configuration is not relevant. As shown in Figure 3, at the start of each iteration, the coordinator schedules up to 𝑁 (e.g., 4) ranks onto available GPUs and begins recording their execution graphs. Each rank runs until it reaches a communication point (e.g., an

5

ESTABLISHING HIGH-FIDELITY GRAPH

5

Bare Graph Rank0 F0, S, F1, S, F2, S, F3... Rank1 R, F1, R, F2, R, F3, R…

Rank0 Rank1

Rank2 R, F1, R, F2, R, B1, S…

Rank2

Rank3 R, F1, B1, S, R, F2, S…

Rank3

…

…

Training Configurations TP=2, PP=4, EP=8, BS=512 …

R B1 R F5 F0 S F1 S F2 S F3 S R R S Slice 0 R F1 S F2 S F3 R B1 R F4 S R B1 R F3 … S R F1 B1 S F2 R B2 …

R F1 S F2 Slice 1

Rank0 Rank1 Rank2 Rank3

…

Calibration

R B1 R F5 … R S F4 … Slice 0 R F1 R R B1 R S F2 S F3 S F3 … R F1 R R B1 R S F2 S Slice 1 R F1 B1 S F2 R B2 …

…

…

Bare Graph Collection

Completed Graph … F0 S F1 S F2 S F3 S

Slice-by-Slice Timing Filling

Ready for Emulate

Inter-slice Timing Calibration

Figure 4. Workflow for generating a complete graph for emulation.

5.3 Stage 2: Timing Filling and Calibration After collecting the bare PrismTrace graph, we know what operations execute and in what order, but not how long they take. This is because graph collection introduces overheads (e.g., context switching) that distort timing. To recover accurate timing with limited GPUs, we divide the training job into subsets, or slices. We execute slices sequentially, using all available GPUs per slice. In each slice, a subset of ranks runs with real computation and communication, while assistant GPUs emulate the remaining ranks. After all slices are executed, each rank has been run at least once as a real rank, yielding accurate slice-level timing. We then combine slice-level graphs and calibrate their timings to construct a globally consistent execution graph. 1. Slice-by-slice timing filling. we execute the workload in slices. In each slice, a subset of ranks runs the original program on available GPUs, while the rest are replayed as virtual ranks on assistant GPUs using the bare graph (capturing what and in what order) to provide realistic communication behavior. This allows each real rank to observe realistic interactions and enables accurate measurement of computation and communication durations. For example, in a 1024-rank workload, a slice may execute ranks 0–7 on sandbox GPUs, while ranks 8–1023 are replayed on assistant GPUs. PrismLLM iterates over such subsets in a round-robin manner (e.g., 0–7, 8–15, . . . ), ensuring every rank is executed once as a real rank. 2. Inter-slice timing calibration. Timing within each slice is locally accurate but not globally aligned. We therefore use dependency information in the execution graph to align timestamps across slices. As shown in Figure 4, consider a pipeline-parallel job (degree 4) split into two slices. Each slice produces its own execution graph (e.g., slice 0 and slice 1) with locally accurate timing. We then calibrate across slices using dependencies: for example, a receive in slice 1 depends on a send in slice 0, so we shift the receive to occur after the send. Propagating such constraints across the graph reconstructs consistent cross-rank timing and removes artifacts from limited-GPU execution. After calibration, we obtain a high-fidelity graph that accurately reflects large-scale training behavior.

all-reduce). The coordinator then checks whether all participants of the collective are active. If not, the collective cannot proceed. For example, in Figure 3, ranks R0, R1, R2, R3, R6, R8, R10, and R11 belong to the same collective, but not all are active simultaneously, causing execution to block. The active rank then stores its communication input tensors in CPU memory, checkpoints its GPU context, and swaps out to free the GPU. Next, the coordinator schedules ranks that can unblock the collective (e.g., R6, R8, R10, and R11), prioritizing such ranks to reduce context switching (Appendix §A). Once all participants reach the collective and their inputs are available, PrismLLM executes the collective on the CPU and produces outputs. These outputs unblock stalled ranks, which are later rescheduled and resume execution from where they paused. This process—execute, block, swap out, gather tensors, resolve communication, and resume—repeats until all ranks complete the iteration. By the end, PrismLLM captures the full execution graph across all ranks, including operations (what) and dependencies (in what order). Timing (how long) is added later during calibration (§5.3). User-defined communication input. While context switching correctly handles value-dependent logic, it can incur non-trivial overhead when GPUs are limited (Appendix §B). For advanced users who have prior knowledge of their workload, PrismLLM provides an alternative: directly specifying the input and output tensors required for communication, bypassing context-switching entirely. These tensors can be obtained from pre-recorded communication data from prior large-scale runs or generated using user-defined rules via our tensor generator. The generator takes user-specified rules to produce the required tensors (examples in Appendix §C), and can be customized as needed. With user-provided tensors, each rank can execute independently using its corresponding inputs on the available GPUs, as shown in the lower part of Figure 3. For example, for a 1024-GPU workload with 4 physical GPUs, PrismLLM executes 256 rounds to cover all ranks, producing the full execution graph (without timing). This can be further reduced by a factor of 1/𝑁 , where 𝑁 is the data-parallel (DP) group size, since execution graphs are identical across DP groups. With this optimization, the number of rounds drops significantly, reducing graph collection time to minutes.

6

HYBRID EMULATION

After collecting the execution graph, we perform emulation to obtain performance metrics. We next describe how the execution graph is replayed (§6.1). We then present how 6

our design minimizes GPU cost during hybrid emulation by mapping virtual ranks onto physical assistant GPUs with low overhead through virtualized initialization (§6.2) and communication pruning (§6.3).

Neighbor Ranks

Sandbox Ranks

Neighbor Ranks Pruned Ranks

6.1 Assistant GPU Workflow for Emulation We execute the real program on sandbox GPUs while replaying many virtual ranks on assistant GPUs, enabling faithful large-scale behavior with only a few GPUs. For virtual ranks, PrismLLM initially instantiates only the subset of them that directly communicate with sandbox GPUs (cost discussed in §6.2). Once initialized, these ranks participate in replay. During replay, virtual ranks do not perform real computation. Instead, they traverse the precollected execution graph using its timing and dependency information, thereby assisting sandbox GPUs in completing communication. The graph contains both computation and communication nodes. When reaching a computation node (e.g., microbatch #4), a virtual rank waits for the recorded duration. When reaching a communication node (e.g., send #3), it performs the actual communication with sandbox ranks. To support this process, PrismLLM manages replayed tensors through a runtime prefetch pipeline. Tensors required for upcoming collective operations are prefetched during idle communication periods to avoid contention. We prioritize loading tensors directly from disk into a GPU buffer pool; if GPU memory is full, tensors are staged in a CPU buffer pool until space becomes available. Together, these buffers act as staging caches. In practice, the GPU buffer pool is typically sufficient. For example, in a 512-GPU job with a single assistant GPU machine, each GPU provides over 100 GB of buffer space for replay (excluding CUDA context usage such as NCCL). Moreover, only tensors involved in communication between sandbox ranks and active virtual ranks need to be managed, significantly reducing buffering requirements. As a result, when a communication node is reached, the required tensor is already resident on the GPU. After communication completes, the tensor is evicted. This design allows a single GPU to efficiently serve multiple virtual ranks and communication groups on demand.

Pruned Ranks

Sandbox Ranks

(a) Ring algorithm. (b) Tree algorithm. Figure 5. Runtime Communication Pruning.

To address this, PrismLLM introduces NCCL group reduction, which reduces overhead by selectively instantiating only the necessary groups and ranks. During group creation, sandbox ranks remain unchanged and are initialized normally. For virtual ranks, we instantiate only the NCCL groups whose members overlap with sandbox ranks. Groups that do not communicate with the sandbox are bypassed at the PyTorch c10d layer and are never instantiated, significantly reducing the total number of active groups. For example, when emulating 2,000 virtual ranks, we reduce the number of active NCCL groups from 8,617 to 82. While reducing the number of NCCL groups helps, instantiating all members within each group remains unnecessary and can be further reduced. As shown in Figure 5, direct communication between the ranks of interest and virtual ranks occurs only among neighboring ranks, depending on the collective algorithm (e.g., ring or tree). Therefore, we can safely prune non-neighboring virtual ranks as long as we preserve numerical correctness. This reduces both initialization cost and runtime communication overhead. Concretely, when all members of a group enter communicator initialization (e.g., dist.new_group() for a DP group), they participate in a TCPStore-based barrier, where each rank increments a counter and waits until all expected ranks have joined. However, since we instantiate only neighboring ranks, the number of active ranks no longer matches the expected group size. To address this, we designate a leader among the assistant ranks to additionally issue requests to TCPStore on behalf of all non-neighbor ranks. This effectively accounts for the missing participants and allows the initialization barrier to complete without modifying the world size or user-level code. Furthermore, during NCCL communicator initialization, each machine typically explores its intra-node topology to identify its logical neighbors. We hijack the NCCL initialization control channel to manage this information exchange. Because the emulated cluster is homogeneous, the Proxy can seamlessly synthesize and exchange the correct intra-node topological data for all pruned ranks. This effectively establishes the correct NCCL topology and communication channels, successfully virtualizing the backend without the sandbox noticing any missing peers.

6.2 Virtual Ranks Initialization Virtual ranks must be placed on assistant GPUs to interact with sandbox GPUs, but hosting thousands of them incurs substantial overhead. This overhead primarily stems from two factors. First, large-scale training creates numerous communication groups (e.g., DP, TP, PP). Each virtual rank requires its own communicator and buffer (e.g., 500 MB per group), which can quickly exhaust GPU memory. Second, initializing thousands of these groups on a few physical GPUs causes severe contention for shared resources like CUDA contexts and driver locks, turning initialization into a partially serialized process that can take hours. 7

Chunk 0 1

vRank 42 ANY ANY

Rank 43 ANY ANY

Rank 44 ANY ANY

vRank 45 ANY ANY

vRank 42 VAL VAL

Chunk 0 1

Rank 43 VAL VAL

Rank 44 VAL VAL

vRank 45 ANY ANY

…

…

VAL VAL VAL VAL

ANY ANY ANY ANY

ANY

ANY

ANY

…

…

ANY

The leftmost vRank prepares the chunks with values pre-subtracted by the sandbox values.

…

… VAL VAL VAL VAL

…

… VAL VAL VAL VAL

…

… 42 43 44 45

…

… ANY ANY ANY ANY

…

… ANY ANY VAL ANY

…

… ANY VAL K-1 ANY

…

… ANY K-1 K-2 ANY

…

… 42 43 44 45

VAL

VAL

VAL

ANY

after the reduction stage is equivalent to that of the original execution, and the values of all other chunks do not affect the outcome. To achieve this, the leftmost virtual rank prepares adjusted values that compensate for missing contributions from pruned ranks. For each chunk, the goal is to ensure that the value received at the owning sandbox rank matches (𝑖 ) datafull . For the chunk owned by rank 44, since the ring path reaches rank 44 via rank 43, the leftmost virtual rank (44) (44) (44) prepares datafull − data43 − data44 , so that as the data

The leftmost vRank prepares all chunks with full values.

(a) Reduce Stage (b) Broadcast Stage Figure 6. Guarantee collective numeric correctness with pruning.

(44) propagates, rank 43 adds back data43 and rank 44 subse6.3 Runtime Communication Pruning (44) quently adds back data44 , reconstructing the correct final Since we remove inactive, non-neighboring ranks, the origi(44) value datafull . Similarly, for other sandbox-owned chunks, nal collective communication pattern is changed at runtime we subtract the contributions of ranks along the remaining (Figure 5). We must therefore ensure that the semantics of path before reaching the owner (including the owner itself). collective operations remain correct under this transformaFor chunks not owned by sandbox ranks, their values do not tion. We modify NCCL’s data transmission logic to guarantee affect the final outcome observed in the sandbox and can numerical correctness. To this end, we modify NCCL to guartherefore be assigned arbitrary values. antee correctness under virtualization. In the broadcast stage, each rank propagates its completed Collective operations can be broadly categorized into two chunk to all other ranks. At this point, the leftmost virtual fundamental types: reduction and broadcast. Many commonly rank needs to provide correct final values for all other chunks used collectives (e.g., all-reduce, reduce-scatter, all-gather) not owned by the sandbox ranks. Since sandbox ranks (e.g., can be decomposed into compositions of these two prim(𝑖 ) ranks 43 and 44) receive the fully reduced values datafull itives, implemented using ring- or tree-based algorithms. Other operations, such as all-to-all and point-to-point send/receive, for their respective chunks during this stage, ensuring correctness for these chunks is sufficient to guarantee overall can be viewed as direct data exchanges. Thus, ensuring corcorrectness from the sandbox’s perspective. rectness for reduction and broadcast is sufficient to cover all Tree-based algorithms follow the same principle; details collective operations. are provided in Appendix §D. We illustrate our approach using ring all-reduce, which consists of both reduction and broadcast stages; the same 7 IMPLEMENTATION principles apply to other collectives. In the reduction stage, each rank splits its data into multiple chunks and is responsiPrismLLM first collects the execution graph using a centralble for aggregating one chunk (e.g., sum, max, or min) from ized program coordinator, and reuses it to enable virtual rank all other ranks. Let data (𝑖𝑗 ) denote the contribution of rank replayers for slice timing collection and emulation. Í (𝑖 ) 𝑗 to chunk 𝑖, and let datafull = 𝑗 data (𝑖𝑗 ) denote the fully We implement a centralized coordinator that orchestrates reduced value. The key property is that each rank only needs all logical ranks during graph collection. It intercepts unto ensure the correctness of its assigned chunk. After one modified training programs and remaps logical ranks onto a full ring round (𝐾 − 1) of communication, each rank holds limited set of GPUs via context switching. We patch the (𝑖 ) PyTorch c10d layer to intercept communication and endatafull for its assigned chunk 𝑖. In the broadcast stage, each able CPU-side collective execution without requiring all parrank then propagates its completed chunk to all other ranks ticipants to be active. A CPU collective executor, built on through another ring round (𝐾 − 1), ensuring that all ranks torch.distributed, processes inputs once all required tenobtain the full result. sors are available and produces outputs that are stored and Figure 6 illustrates how PrismLLM preserves collective reused when ranks resume, ensuring correctness without correctness under rank pruning. The example includes neighfull concurrency. boring virtual ranks (e.g., vRank 42 and vRank 45) and sandThe replayer serves both slice-level timing collection and box ranks (e.g., Rank 43 and Rank 44); other pruned nonfinal emulation. The key difference lies in how the execution neighboring ranks are omitted for clarity. In the reduction graph is used. During slice execution, virtual ranks replay stage, we leverage the key property of ring all-reduce: each only the graph structure without requiring accurate timing, rank is responsible only for its assigned chunk. Therefore, as this does not affect the timing of real ranks within each after pruning, we only need to ensure that sandbox ranks slice; any drift is corrected during global calibration. In concompute their assigned chunks correctly. Once the ranks of trast, during final emulation, virtual ranks execute the fully interest (e.g., ranks 43 and 44) obtain the correct values for calibrated graph with complete timing information. their chunks, i.e., data (43) and data (44) , the resulting state full

full 8

memory is measured using max_memory_allocated in PyTorch. Each baseline includes at least 10 warm-up iterations, followed by monitoring peak memory over 3–10 iterations. For MoE models, memory usage varies across devices due to uneven expert workloads after dispatch. We evaluate two cases: (1) Balanced, using Megatron Core to ensure uniform distribution, and (2) Imbalanced, where we inject precomputed skewed logits via our MoE mock router (Appendix §F) to create non-uniform workloads. Figure 8 presents the emulation results of PrismLLM. For peak memory allocation, PrismLLM demonstrates exceptional accuracy, with prediction errors consistently below 0.01%. Furthermore, PrismLLM successfully reproduces outof-memory (OOM) errors encountered in the baseline experiments. PrismLLM achieves high-fidelity memory estimation by accounting for consumption across the entire software stack, ensuring that its peak memory predictions are both comprehensive and accurate.

We also extend NCCL to support communication pruning. We first identify whether each comm_id involves communication with neighbors of sandbox ranks and pass this information to NCCL. During execution, if a communication corresponds to a non-neighbor comm_id (i.e., not on a path to sandbox ranks), NCCL skips the data transfer and instead generates completion metadata to satisfy the collective, reducing overhead while preserving correctness.

8

EVALUATION

We evaluate PrismLLM across four dimensions. First, end-toend results (§8.1) show that PrismLLM achieves high-fidelity predictions for both iteration time and memory allocation, with average errors of 0.58% and below 0.01%, respectively. Second, we evaluate emulation efficiency and scalability across diverse target configurations (§8.2). Third, we conduct microbenchmarks to assess execution fidelity and quantify the benefits of our key optimizations—bootstrap acceleration 8.2 System Efficiency via virtualized initialization, transmission reduction via comTo evaluate the system efficiency of PrismLLM, we analyze munication pruning, and timing alignment via inter-slice calits physical resource requirements and end-to-end emulation ibration (§8.3). Finally, we compare PrismLLM with state-oftime across various scales. Driven by the common practice of the-art simulators, including Phantora [19] and SimAI [27], expanding the Pipeline Parallelism (PP) degree to accommoshowing that PrismLLM achieves higher estimation accuracy date larger model states, we scale our target workload from for complex training workloads (§8.4). 512 to 8,192 GPUs. Specifically, we proportionally increase Experimental setup. We conduct experiments on a 2,048the PP size from 4 to 64 while holding the DP size constant. GPU testbed (8 GPUs per node) using the open-source MegatronTo prevent linear time overhead when emulating more LM implementation of Qwen 3 MoE pretraining. We evaluate PP stages, PrismLLM parallelizes emulation by scaling Asthree model configurations spanning a wide range of scales: sistant Nodes proportionally to the target scale (red line in 235B-A22B (M.1), 503B-A20B (M.2), and 1.01T-A43B (M.3). Figure 9). A strict 1:1 ratio between Assistant and Sandbox To cover diverse parallelization strategies under memory Nodes enables independent, concurrent profiling of pipeline constraints, we vary tensor parallel (TP), pipeline parallel stages. For example, a 512-GPU target (PP=4) requires 4 se(PP), and expert parallel (EP) sizes, along with gradient accuquential steps on 1 Assistant Node. Scaling to 1,024 GPUs mulation steps. We evaluate four configurations (S.A–S.D). (PP=8) with 2 Assistant Nodes allows concurrent execution, Detailed settings are provided in Appendix §E. keeping the workload at 4 steps and the overall emulation 8.1 End-to-end Prediction Accuracy time comparable to the baseline. This strategy keeps the end-to-end emulation time highly End-to-end iteration time estimation. We evaluate three stable (under 80 minutes) up to 8,192 GPUs. Without proporcluster scales (512, 1,024, and 2,048 GPUs), each with three tional scaling, a 1,024-GPU ablation with 1 Assistant Node model configurations and two parallelization strategies. Itera(second bar) must process all 8 stages sequentially, nearly tion time is defined as the elapsed time reported by Megatrondoubling the time (125 minutes). Crucially, PrismLLM mainLM for one step, including both computation and commutains massive efficiency: emulating an 8,192-GPU cluster nication. Each baseline runs 85 iterations for stability, and requires only 32 physical nodes in total (16 Assistant, 16 PrismLLM uses identical configurations and training scripts. Sandbox), utilizing < 0.4% of the target GPUs and achieving Figure 7 compares PrismLLM and SimAI against the base> 99% resource savings. In the implementation of hybrid emline. PrismLLM achieves high accuracy across all scales, with ulation, PrismLLM mitigates the overhead associated with an average error of 0.58% and a maximum error of 1.98%, redundant virtual rank bootstrapping and inter-rank comdemonstrating its ability to accurately predict end-to-end munication, thereby ensuring high system efficiency. iteration time under diverse parallelization strategies. The

8.3 Micro benchmarks Execution fidelity. To validate the execution accuracy inside PrismLLM, we compare its PyTorch Profiler traces against training baselines. We perform a fine-grained, kernel-level comparison using the absolute deviation in kernel duration.

use of actual GPUs for task execution, coupled with the construction of high-fidelity dependency graphs, empowers PrismLLM to achieve highly accurate latency estimations. End-to-end peak memory allocation estimation. We evaluate peak memory under the same three scales. Peak 9

Iteration Time (s)

Baseline

Scale: 512 GPUs

PrismLLM

Scale: 1024 GPUs

Scale: 2048 GPUs +0.17% 12 -0.40% 12 +0.48% 12 9 9 9 +0.48% -0.24% -0.02% -0.92% -0.15% +0.26% +0.60% -1.98% 6 -0.62% -0.52% -0.12% 6 -0.37% +0.81% +1.54% 6 +0.85% 3 3 3 0 235B 235B 503B 503B 1.01T 1.01T 0 235B 235B 503B 503B 1.01T 1.01T 0 235B 235B 503B 503B 1.01T 1.01T (S.A) (S.B) (S.A) (S.B) (S.C) (S.D) (S.A) (S.B) (S.A) (S.B) (S.C) (S.D) (S.A) (S.B) (S.A) (S.B) (S.C) (S.D) Figure 7. End-to-end iteration time estimation results.

Normalized Memory

MoE Balance (Baseline)

Scale: 512 GPUs

0.9 0.6 0.3 0

MoE Balance (PrismLLM)

0.9 0.6 0.3 OOM OOM OOM 0 235B 235B 503B 503B 1.01T 1.01T (S.A) (S.B) (S.A) (S.B) (S.C) (S.D)

MoE Imbalance (Baseline)

Scale: 1024 GPUs

0.9 0.6 0.3 OOM OOM OOM 0 235B 235B 503B 503B 1.01T 1.01T (S.A) (S.B) (S.A) (S.B) (S.C) (S.D)

MoE Imbalance (PrismLLM)

Scale: 2048 GPUs

OOM

235B (S.A)

OOM

235B (S.B)

503B (S.A)

503B (S.B)

OOM

1.01T 1.01T (S.C) (S.D)

Time Fill & Calibration

120 80 40 0

16 12 8 4 0

Assistant Nodes

Elapsed Time (min)

Emulate

512 10241024204840968192 Scale

Figure 9. Emulation time breakdown.

Dev. from Median (µs)

For imbalanced cases, the MoE mock router (Appendix §F) metrics are br_min: 0.71, br_max: 2.16, br_avg: 1.48, br_std: 0.37, br_med: 1.38, br_skew: 0.90. Figure 8. End-to-end peak memory allocation estimation results. All estimation errors are consistently below 0.01%.

Natural Variance

Compute Kernels

PrismLLM

NCCL Kernels

10 0 10 235B 235B 503B 503B 1.01T 1.01T (S.A) (S.B) (S.A) (S.B) (S.C) (S.D)

235B 235B 503B 503B 1.01T 1.01T (S.A) (S.B) (S.A) (S.B) (S.C) (S.D)

Figure 10. Kernel duration deviation from the baseline median under the 1024-rank scale.

For the training baseline, we sample 32 ranks within the same DP group and calculate the median execution duration of each kernel. The inherent hardware jitter—the deviation of actual execution times from this median—is represented by the Natural Variance boxes. Then, we compare the emulated kernel durations against the corresponding baseline medians, with the resulting errors shown by the PrismLLM boxes. As shown in Figure 10, for all configurations under the scale of 1024 ranks, the emulated variance falls within the range of natural hardware variance. Beyond duration, we evaluate scheduling fidelity via relative kernel start times (normalized by total iteration time). Figure 12 visualizes the 1024-rank 503B (S.B) case, plotting the natural variance across 32 baseline ranks (gray) against PrismLLM’s deviation from the baseline median (blue). PrismLLM tightly tracks the median, bounded within the ±0.5% natural hardware jitter. Across all configurations, the maximum observed start time error is 3.64%. By leveraging physical GPUs to execute training tasks, PrismLLM minimizes emulation errors of kernels. Virtual rank bootstrap optimization. To enable efficient large-scale emulation, PrismLLM features a customized bootstrapping mechanism. We evaluate its performance by measuring the time and memory required to initialize and execute a global barrier operation across varying cluster scales (64 to 8192 ranks). The testbed consists of two 8-GPU nodes: one acts as the sandbox (running one real rank per GPU),

while the other hosts all virtualized ranks. We compare PrismLLM against a Vanilla baseline that uses standard NCCL (via shared NCCL_HOSTID), where each virtual rank maintains an independent process and CUDA context. As shown in Figures 11a and 11b, the Vanilla approach incurs linearly scaling overheads. At merely 512 ranks, initializing CUDA contexts, establishing communication channels, and allocating NCCL buffers consumes 544.5 GB of CPU memory and 103.7 GB of GPU memory. Consequently, the bootstrap process takes 756.0 seconds to complete, and the baseline inevitably encounters OOM errors at scales ≥ 1024 ranks. In contrast, PrismLLM maintains a strictly constant resource footprint regardless of the logical cluster size. By strategically allocating resources only for "active" virtual ranks (i.e., direct topological neighbors of the sandbox nodes), PrismLLM eliminates massive redundant contexts and buffers. Consequently, PrismLLM successfully bootstraps an 8192-rank emulation in just 44.3 seconds, occupying a mere 5.4 GB of GPU memory per physical device. Virtual rank communication pruning. While bootstrap optimization minimizes initialization overhead, PrismLLM must also mitigate runtime resource contention during collective operations. Using the same two-node testbed, we evaluate AllReduce execution. In the Vanilla solution, virtual ranks blindly replay standard NCCL kernels, intensely 10

102 101 100

OOM

OOM

OOM

OOM

64 128 256 512 1024 2048 4096 8192 Scale

102

CUDA CTX Vanilla PrismLLM NCCL Buffer Vanilla PrismLLM

101 100

OOM

OOM

OOM

OOM

64 128 256 512 1024 2048 4096 8192 Scale

Operation Time (s)

Vanilla PrismLLM

GPU Memory Usage (GB)

CPU Memory Usage (GB)

103

103 OOM OOM OOM OOM

Vanilla PrismLLM

102 101

64 128 256 512 1024 2048 4096 8192

Scale

0

20000 40000 60000 80000

102 101 100 10 1

Kernel Index

Figure 12. Kernel launch deviation.

16M

128M 1G 8G Message Size

32G

Figure 13. Latency overhead.

contending with sandbox ranks for Streaming Multiprocessors and memory/PCIe bandwidth, thereby skewing execution timing. To ensure high fidelity, PrismLLM aggressively prunes redundant data transmissions among virtual ranks, preserving only essential cross-node RDMA traffic. Figure 13 shows the transmission latency across varying message sizes at a 32-rank scale. Without pruning, the Vanilla approach suffers up to a 148× latency inflation. By applying our pruning techniques, PrismLLM closely tracks the physical baseline, yielding a maximum error of 0.2% against real executions. Further evaluations up to 128 ranks are in Appendix §G. Inter-slice calibration. To demonstrate the effectiveness of inter-slice timing calibration, we chose the 235B model in strategy S.B and replayed the graph before and after interslice calibration. The emulated iteration time before interslice calibration quickly drops from 5.7s to 5.13s, demonstrating that without proper alignment, the emulation error can easily go higher than 10%. 8.4

4 0

4B seqlen 4096

4B seqlen 1024

85%

Memory

+0.

7%

SimAI

-0.3

2%

2

-0.6

6

Phantora

0.9

Normalized Memory

-0.50%

103

PrismLLM

0.6 -79 .93%

-0.25%

Baseline

-79 .37%

0.00%

PrismLLM

-79 .39%

0.25%

Vanilla

-0.4 5% -64 .29%

0.50%

Baseline

104

-1.1 -12 5% .64%

PrismLLM

Iteration Time (s)

Natural Variance

0.75%

Transmission Latency (ms)

Relative Start Time Dev. (%)

(a) Overall CPU memory overhead. (b) GPU memory overhead per rank. (c) Completion time of one barrier. Figure 11. Memory usage reduction and bootstrap acceleration of large-scale emulation in PrismLLM.

235B (S.A) 235B (S.A) 235B (S.A) scale scale scale 512 1024 2048

0.3 0

Figure 14. PrismLLM compared with Simulators.

nature of collective kernels; and (2) Phantora launches multiple mock processes on CPU to represent individual training ranks, advancing the simulation via inter-process communication with a central simulator. Due to the CPU contention among multiple processes, Phantora could not accurately capture the impact of CPU-side Python interpreter latency on the overall execution duration. As a result, Phantora’s defaults to omitting these latency components, leading to a systematic underestimation of the total iteration time. SimAI. To further evaluate the performance of simulators in large-scale MoE training scenarios, we feed the model configurations and parallelization strategies into SimAI [27] and compare its results against PrismLLM. Notably, SimAI does not support memory allocation simulation and is therefore excluded from the memory-related comparison. As shown in Figure 14, we present the results for the 235B (M.1) model with strategy S.A at different scales, while comprehensive results for other configurations are deferred to Appendix §H. SimAI fails to achieve acceptable fidelity, with an average error of 77.2%. We attribute this discrepancy to two primary factors. First, while SimAI relies on workload files to describe the training process, the design of workload file fail to capture the inter-rank dependencies and asymmetric behaviors inherent in pipeline model parallel (PP), omitting the communication bubbles associated with PP. Second, SimAI neglects the unique characteristics of MoE models compared to dense ones, ignoring the computational overheads of critical MoE operations such as gating, permute, dispatch, and combine. These findings further underscore that traditional simulators often require substantial manual effort to maintain and update simulation models for evolving architectures.

Compared with Simulators

Phantora. Due to Phantora’s lack of support for pipeline model parallel (PP) and MoE models, we evaluate PrismLLM against Phantora [19] on small-scale tasks. We assess their accuracy in estimating iteration time and peak memory allocation under identical model and parallelization configurations. The experimental setup is tailored to a 32-GPU testbed using Qwen 3 4B model configuration, with the TP size set to 4 to ensure a fair comparison. As shown in Figure 14, PrismLLM outperforms Phantora in both metrics. Notably, While Phantora shows a 12% error in typical cases, its prediction fidelity drops significantly as the sequence length decreases—with the error surging to 64% as the GPU-side computational pressure is decreased. This discrepancy primarily stems from two factors: (1) Phantora approximates actual kernel duration using raw network bandwidth, failing to account for the complex, multi-step

9

USE CASES AND LESSONS LEARNED

Tuning training configurations. PrismLLM enables zeroeffort tuning of diverse training configurations. Finding the 11

Table 1. Emulation accuracy under different training configurations. Optimization Baseline Flash Attention Off P2P Overlap Off Offload Optimizer Recompute

Iteration Time (ms)

Peak Memory (GB)

Baseline

PrismLLM

Baseline

PrismLLM

5619.8 (+289.1/-82.9) 7737.4 (+238.8/-81.6) 5706.7 (+377.0/-90.5) 11979.2 (+513.9/-237.6) 7074.9 (+372.5/-79.1)

5702.9 7698.9 5818.9 11614.4 7205.0

133.76 135.22 133.74 112.07 98.13

133.76 135.22 133.74 112.07 98.13

optimal strategy, such as balancing batch size, sequence length, parallel dimensions, and memory/computation tradeoffs, is critical for maximizing infrastructure utilization. PrismLLM seamlessly evaluates these configurations, providing high-fidelity emulation without requiring additional engineering overhead. As shown in Table 1, we evaluate four common optimization techniques: CPU offloading, recomputation, P2P overlapping, and switching computation libraries. PrismLLM accurately captures the corresponding iteration time and memory usage for each. MoE memory estimation and balance control. Estimating memory for MoE models is notoriously difficult because token routing is dynamic, input-dependent, and shifts across training stages. Early training exhibits severe routing imbalance, which can increase the end-to-end iteration time by over 5% and cause memory usage to fluctuate by up to 20 GB. Consequently, aggressive memory configurations risk OOM errors, while conservative ones waste GPU resources. PrismLLM solves this by allowing engineers to plug in a mock router and replay production-observed imbalance features (e.g., minimum/maximum token proportions per device). By reverse-computing these features, PrismLLM accurately estimates dynamic memory usage, enabling engineers to adaptively select optimal training configurations for different training stages. Targeted cluster health checks. Production incidents caused by gray failures are notoriously hard to pinpoint. Traditional benchmark suites often fail to cover the massive search space of complex software stacks and corner cases. PrismLLM transforms pre-training health checks by executing the exact target workload on isolated replaying units of the cluster before a full-scale launch. This targeted approach exposes both explicit and gray failures. In one example, one of our GPU servers has a thermal issue, the GPU’s graphics frequency is down-clocked to 900MHz due to thermal throttling, the entire training job’s training iteration time is slowed down by 14% (from 5.63s to 6.38s). Small-scaled health checks are not able to reproduce this problem since they cannot push the machine hard enough to hit the thermal issue. While PrismLLM is able to perform pair-wise checking and reproduce the issue. The average training iteration replayed by PrismLLM increases from 5.70s to 6.40s. The entire training job returns to normal after bringing the specific GPU offline. Fine-grained workload profiling and monitoring. Many debugging and fine-grained workload profiling tools are prohibited in the online workload due to their high performance overhead, pressure to the logging system, or potential bugs.

With PrismLLM, we can safely inspect the training stack, deploy debugging tools, or test the compatibility of a new enhancement without deploying in the production. All of our evaluation results are collected using the vanilla monitoring and tracing tools, such as nvidia-smi, PyTorch Profiler, and NVidia Nsight System without any modification. Partial graph re-alignment. Note that many enhancements focus on improving the performance of a single kernel, or local communication-computation overlapping. Under such cases, we don’t need to go through bare-graph generation once again since there is no structural change, we only need to re-calibrate the intra- and inter-slice timing. Optimization planning. It’s non-trivial to estimate the gain of an optimization before implementing it. PrismLLM allows the developer to plan before implementation by creating a fake GPU kernel that spins for the desired and optimized duration, and then evaluating the end-to-end performance.

10

RELATED WORK AND DISCUSSION

Simulating LLM training performance has become an active area of research. Existing simulation frameworks [1, 9, 10, 13, 19, 21, 27, 28, 33] estimate performance by profiling or collecting execution traces and then performing offline analysis. While these approaches provide useful estimates, they are often limited in flexibility and cannot fully capture complex system behavior. In particular, simulation-based approaches struggle to reproduce real system issues (e.g., bugs or performance anomalies) because they abstract away many runtime interactions. For example, Phantora [19] allows users to execute the original training program with replaced CUDA and NCCL backends. However, because these components are simulated rather than real, the execution is not fully faithful. NCCL behavior is complex—e.g., buffer movement and the overlapping of communication and computation—which further reduces simulation accuracy. Moreover, these approaches cannot faithfully capture system-level memory dynamics (e.g., fragmentation and allocator behavior below CUDA), which depend on actual runtime execution and closed-source behaviors that cannot be easily inferred. Limitation. PrismLLM highly relies on the assumption that the training program is runnable on the target GPU platform, which means we cannot conduct forward planning, i.e., predict the performance on an unknown hardware before its release. We also rely on the assumption that the GPU’s workload stays periodic across iterations, which means for distillation workloads where the number of active teacher models varies across iterations, PrismLLM can only emulate one representative workload.

11

CONCLUSION

We present PrismLLM, a system that enables faithful largescale training performance estimation and debugging using limited hardware. By constructing high-fidelity execution graphs and replaying them at low cost, PrismLLM allows a 12

small number of GPUs to emulate deployment-scale behavior while preserving both performance characteristics and program semantics.

13

REFERENCES

[8] Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshorn, Aobo Yang, Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurelien Rodriguez, Austen Gregerson, Ava Spataru, Baptiste Roziere, Bethany Biron, Binh Tang, Bobbie Chern, Charlotte Caucheteux, Chaya Nayak, Chloe Bi, Chris Marra, Chris McConnell, Christian Keller, Christophe Touret, Chunyang Wu, Corinne Wong, Cristian Canton Ferrer, Cyrus Nikolaidis, Damien Allonsius, Daniel Song, Danielle Pintz, Danny Livshits, Danny Wyatt, David Esiobu, Dhruv Choudhary, Dhruv Mahajan, Diego Garcia-Olano, Diego Perino, Dieuwke Hupkes, Egor Lakomkin, Ehab AlBadawy, Elina Lobanova, Emily Dinan, Eric Michael Smith, Filip Radenovic, Francisco Guzmán, Frank Zhang, Gabriel Synnaeve, Gabrielle Lee, Georgia Lewis Anderson, Govind Thattai, Graeme Nail, Gregoire Mialon, Guan Pang, Guillem Cucurell, Hailey Nguyen, Hannah Korevaar, Hu Xu, Hugo Touvron, Iliyan Zarov, Imanol Arrieta Ibarra, Isabel Kloumann, Ishan Misra, Ivan Evtimov, Jack Zhang, Jade Copet, Jaewon Lee, Jan Geffert, Jana Vranes, Jason Park, Jay Mahadeokar, Jeet Shah, Jelmer van der Linde, Jennifer Billock, Jenny Hong, Jenya Lee, Jeremy Fu, Jianfeng Chi, Jianyu Huang, Jiawen Liu, Jie Wang, Jiecao Yu, Joanna Bitton, Joe Spisak, Jongsoo Park, Joseph Rocca, Joshua Johnstun, Joshua Saxe, Junteng Jia, Kalyan Vasuden Alwala, Karthik Prasad, Kartikeya Upasani, Kate Plawiak, Ke Li, Kenneth Heafield, Kevin Stone, Khalid El-Arini, Krithika Iyer, Kshitiz Malik, Kuenley Chiu, Kunal Bhalla, Kushal Lakhotia, Lauren Rantala-Yeary, Laurens van der Maaten, Lawrence Chen, Liang Tan, Liz Jenkins, Louis Martin, Lovish Madaan, Lubo Malo, Lukas Blecher, Lukas Landzaat, Luke de Oliveira, Madeline Muzzi, Mahesh Pasupuleti, Mannat Singh, Manohar Paluri, Marcin Kardas, Maria Tsimpoukelli, Mathew Oldham, Mathieu Rita, Maya Pavlova, Melanie Kambadur, Mike Lewis, Min Si, Mitesh Kumar Singh, Mona Hassan, Naman Goyal, Narjes Torabi, Nikolay Bashlykov, Nikolay Bogoychev, Niladri Chatterji, Ning Zhang, Olivier Duchenne, Onur Çelebi, Patrick Alrassy, Pengchuan Zhang, Pengwei Li, Petar Vasic, Peter Weng, Prajjwal Bhargava, Pratik Dubal, Praveen Krishnan, Punit Singh Koura, Puxin Xu, Qing He, Qingxiao Dong, Ragavan Srinivasan, Raj Ganapathy, Ramon Calderer, Ricardo Silveira Cabral, Robert Stojnic, Roberta Raileanu, Rohan Maheswari, Rohit Girdhar, Rohit Patel, Romain Sauvestre, Ronnie Polidoro, Roshan Sumbaly, Ross Taylor, Ruan Silva, Rui Hou, Rui Wang, Saghar Hosseini, Sahana Chennabasappa, Sanjay Singh, Sean Bell, Seohyun Sonia Kim, Sergey Edunov, Shaoliang Nie, Sharan Narang, Sharath Raparthy, Sheng Shen, Shengye Wan, Shruti Bhosale, Shun Zhang, Simon Vandenhende, Soumya Batra, Spencer Whitman, Sten Sootla, Stephane Collot, Suchin Gururangan, Sydney Borodinsky, Tamar Herman, Tara Fowler, Tarek Sheasha, Thomas Georgiou, Thomas Scialom, Tobias Speckbacher, Todor Mihaylov, Tong Xiao, Ujjwal Karn, Vedanuj Goswami, Vibhor Gupta, Vignesh Ramanathan, Viktor Kerkez, Vincent Gonguet, Virginie Do, Vish Vogeti, Vítor Albiero, Vladan Petrovic, Weiwei Chu, Wenhan Xiong, Wenyin Fu, Whitney Meers, Xavier Martinet, Xiaodong Wang, Xiaofang Wang, Xiaoqing Ellen Tan, Xide Xia, Xinfeng Xie, Xuchao Jia, Xuewei Wang, Yaelle Goldschlag, Yashesh Gaur, Yasmine Babaei, Yi Wen, Yiwen Song, Yuchen Zhang, Yue Li, Yuning Mao, Zacharie Delpierre Coudert, Zheng Yan, Zhengxing Chen, Zoe Papakipos, Aaditya Singh, Aayushi Srivastava, Abha Jain, Adam Kelsey, Adam Shajnfeld, Adithya Gangidi, Adolfo Victoria, Ahuva Goldstand, Ajay Menon, Ajay Sharma, Alex Boesenberg, Alexei Baevski, Allie Feinstein, Amanda Kallet, Amit Sangani, Amos Teo, Anam Yunus, Andrei Lupu, Andres Alvarado, Andrew Caples, Andrew Gu, Andrew Ho, Andrew Poulton, Andrew Ryan, Ankit Ramchandani, Annie Dong, Annie Franco, Anuj Goyal, Aparajita Saraf, Arkabandhu Chowdhury, Ashley Gabriel, Ashwin Bharambe, Assaf Eisenman, Azadeh Yazdan, Beau James, Ben Maurer, Benjamin Leonhardi, Bernie Huang, Beth Loyd, Beto De Paola, Bhargavi Paranjape, Bing Liu, Bo Wu, Boyu

[1] Jehyeon Bang, Yujeong Choi, Myeongwoo Kim, Yongdeok Kim, and Minsoo Rhu. 2024. vTrain: A Simulation Framework for Evaluating Cost-Effective and Compute-Optimal Large Language Model Training. In Proceedings of the 2024 57th IEEE/ACM International Symposium on Microarchitecture (MICRO ’24). IEEE Press, 153–167. https://doi.org/ 10.1109/MICRO61859.2024.00021 [2] Li-Wen Chang, Wenlei Bao, Qi Hou, Chengquan Jiang, Ningxin Zheng, Yinmin Zhong, Xuanrun Zhang, Zuquan Song, Ziheng Jiang, Haibin Lin, Xin Jin, and Xin Liu. 2024. FLUX: Fast Software-based Communication Overlap On GPUs Through Kernel Fusion. (2024). arXiv:cs.LG/2406.06858 [3] CRIU Project Developers. 2026. Github - CRIU: Checkpoint/Restore In Userspace. https://github.com/checkpoint-restore/criu. (2026). [4] Tri Dao. 2023. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. (2023). arXiv:cs.LG/2307.08691 https://arxiv.org/abs/2307.08691 [5] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. 2022. FLASHATTENTION: fast and memory-efficient exact attention with IO-awareness. In Proceedings of the 36th International Conference on Neural Information Processing Systems (NIPS ’22). Curran Associates Inc., Red Hook, NY, USA, Article 1189, 16 pages. [6] DeepSeek-AI, Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, Damai Dai, Daya Guo, Dejian Yang, Deli Chen, Dongjie Ji, Erhang Li, Fangyun Lin, Fucong Dai, Fuli Luo, Guangbo Hao, Guanting Chen, Guowei Li, H. Zhang, Han Bao, Hanwei Xu, Haocheng Wang, Haowei Zhang, Honghui Ding, Huajian Xin, Huazuo Gao, Hui Li, Hui Qu, J. L. Cai, Jian Liang, Jianzhong Guo, Jiaqi Ni, Jiashi Li, Jiawei Wang, Jin Chen, Jingchang Chen, Jingyang Yuan, Junjie Qiu, Junlong Li, Junxiao Song, Kai Dong, Kai Hu, Kaige Gao, Kang Guan, Kexin Huang, Kuai Yu, Lean Wang, Lecong Zhang, Lei Xu, Leyi Xia, Liang Zhao, Litong Wang, Liyue Zhang, Meng Li, Miaojun Wang, Mingchuan Zhang, Minghua Zhang, Minghui Tang, Mingming Li, Ning Tian, Panpan Huang, Peiyi Wang, Peng Zhang, Qiancheng Wang, Qihao Zhu, Qinyu Chen, Qiushi Du, R. J. Chen, R. L. Jin, Ruiqi Ge, Ruisong Zhang, Ruizhe Pan, Runji Wang, Runxin Xu, Ruoyu Zhang, Ruyi Chen, S. S. Li, Shanghao Lu, Shangyan Zhou, Shanhuang Chen, Shaoqing Wu, Shengfeng Ye, Shengfeng Ye, Shirong Ma, Shiyu Wang, Shuang Zhou, Shuiping Yu, Shunfeng Zhou, Shuting Pan, T. Wang, Tao Yun, Tian Pei, Tianyu Sun, W. L. Xiao, Wangding Zeng, Wanjia Zhao, Wei An, Wen Liu, Wenfeng Liang, Wenjun Gao, Wenqin Yu, Wentao Zhang, X. Q. Li, Xiangyue Jin, Xianzu Wang, Xiao Bi, Xiaodong Liu, Xiaohan Wang, Xiaojin Shen, Xiaokang Chen, Xiaokang Zhang, Xiaosha Chen, Xiaotao Nie, Xiaowen Sun, Xiaoxiang Wang, Xin Cheng, Xin Liu, Xin Xie, Xingchao Liu, Xingkai Yu, Xinnan Song, Xinxia Shan, Xinyi Zhou, Xinyu Yang, Xinyuan Li, Xuecheng Su, Xuheng Lin, Y. K. Li, Y. Q. Wang, Y. X. Wei, Y. X. Zhu, Yang Zhang, Yanhong Xu, Yanhong Xu, Yanping Huang, Yao Li, Yao Zhao, Yaofeng Sun, Yaohui Li, Yaohui Wang, Yi Yu, Yi Zheng, Yichao Zhang, Yifan Shi, Yiliang Xiong, Ying He, Ying Tang, Yishi Piao, Yisong Wang, Yixuan Tan, Yiyang Ma, Yiyuan Liu, Yongqiang Guo, Yu Wu, Yuan Ou, Yuchen Zhu, Yuduan Wang, Yue Gong, Yuheng Zou, Yujia He, Yukun Zha, Yunfan Xiong, Yunxian Ma, Yuting Yan, Yuxiang Luo, Yuxiang You, Yuxuan Liu, Yuyang Zhou, Z. F. Wu, Z. Z. Ren, Zehui Ren, Zhangli Sha, Zhe Fu, Zhean Xu, Zhen Huang, Zhen Zhang, Zhenda Xie, Zhengyan Zhang, Zhewen Hao, Zhibin Gou, Zhicheng Ma, Zhigang Yan, Zhihong Shao, Zhipeng Xu, Zhiyu Wu, Zhongyu Zhang, Zhuoshu Li, Zihui Gu, Zijia Zhu, Zijun Liu, Zilin Li, Ziwei Xie, Ziyang Song, Ziyi Gao, and Zizheng Pan. 2025. DeepSeek-V3 Technical Report. (2025). arXiv:cs.CL/2412.19437 https://arxiv.org/abs/2412.19437 [7] Epoch AI. 2025. "Data on AI Models". (7 2025). https://epoch.ai/data/ ai-models/ Accessed: 13 Mar 2026. 14

Ni, Braden Hancock, Bram Wasti, Brandon Spence, Brani Stojkovic, Brian Gamido, Britt Montalvo, Carl Parker, Carly Burton, Catalina Mejia, Ce Liu, Changhan Wang, Changkyu Kim, Chao Zhou, Chester Hu, Ching-Hsiang Chu, Chris Cai, Chris Tindal, Christoph Feichtenhofer, Cynthia Gao, Damon Civin, Dana Beaty, Daniel Kreymer, Daniel Li, David Adkins, David Xu, Davide Testuggine, Delia David, Devi Parikh, Diana Liskovich, Didem Foss, Dingkang Wang, Duc Le, Dustin Holland, Edward Dowling, Eissa Jamil, Elaine Montgomery, Eleonora Presani, Emily Hahn, Emily Wood, Eric-Tuan Le, Erik Brinkman, Esteban Arcaute, Evan Dunbar, Evan Smothers, Fei Sun, Felix Kreuk, Feng Tian, Filippos Kokkinos, Firat Ozgenel, Francesco Caggioni, Frank Kanayet, Frank Seide, Gabriela Medina Florez, Gabriella Schwarz, Gada Badeer, Georgia Swee, Gil Halpern, Grant Herman, Grigory Sizov, Guangyi, Zhang, Guna Lakshminarayanan, Hakan Inan, Hamid Shojanazeri, Han Zou, Hannah Wang, Hanwen Zha, Haroun Habeeb, Harrison Rudolph, Helen Suk, Henry Aspegren, Hunter Goldman, Hongyuan Zhan, Ibrahim Damlaj, Igor Molybog, Igor Tufanov, Ilias Leontiadis, Irina-Elena Veliche, Itai Gat, Jake Weissman, James Geboski, James Kohli, Janice Lam, Japhet Asher, Jean-Baptiste Gaya, Jeff Marcus, Jeff Tang, Jennifer Chan, Jenny Zhen, Jeremy Reizenstein, Jeremy Teboul, Jessica Zhong, Jian Jin, Jingyi Yang, Joe Cummings, Jon Carvill, Jon Shepard, Jonathan McPhie, Jonathan Torres, Josh Ginsburg, Junjie Wang, Kai Wu, Kam Hou U, Karan Saxena, Kartikay Khandelwal, Katayoun Zand, Kathy Matosich, Kaushik Veeraraghavan, Kelly Michelena, Keqian Li, Kiran Jagadeesh, Kun Huang, Kunal Chawla, Kyle Huang, Lailin Chen, Lakshya Garg, Lavender A, Leandro Silva, Lee Bell, Lei Zhang, Liangpeng Guo, Licheng Yu, Liron Moshkovich, Luca Wehrstedt, Madian Khabsa, Manav Avalani, Manish Bhatt, Martynas Mankus, Matan Hasson, Matthew Lennie, Matthias Reso, Maxim Groshev, Maxim Naumov, Maya Lathi, Meghan Keneally, Miao Liu, Michael L. Seltzer, Michal Valko, Michelle Restrepo, Mihir Patel, Mik Vyatskov, Mikayel Samvelyan, Mike Clark, Mike Macey, Mike Wang, Miquel Jubert Hermoso, Mo Metanat, Mohammad Rastegari, Munish Bansal, Nandhini Santhanam, Natascha Parks, Natasha White, Navyata Bawa, Nayan Singhal, Nick Egebo, Nicolas Usunier, Nikhil Mehta, Nikolay Pavlovich Laptev, Ning Dong, Norman Cheng, Oleg Chernoguz, Olivia Hart, Omkar Salpekar, Ozlem Kalinli, Parkin Kent, Parth Parekh, Paul Saab, Pavan Balaji, Pedro Rittner, Philip Bontrager, Pierre Roux, Piotr Dollar, Polina Zvyagina, Prashant Ratanchandani, Pritish Yuvraj, Qian Liang, Rachad Alao, Rachel Rodriguez, Rafi Ayub, Raghotham Murthy, Raghu Nayani, Rahul Mitra, Rangaprabhu Parthasarathy, Raymond Li, Rebekkah Hogan, Robin Battey, Rocky Wang, Russ Howes, Ruty Rinott, Sachin Mehta, Sachin Siby, Sai Jayesh Bondu, Samyak Datta, Sara Chugh, Sara Hunt, Sargun Dhillon, Sasha Sidorov, Satadru Pan, Saurabh Mahajan, Saurabh Verma, Seiji Yamamoto, Sharadh Ramaswamy, Shaun Lindsay, Shaun Lindsay, Sheng Feng, Shenghao Lin, Shengxin Cindy Zha, Shishir Patil, Shiva Shankar, Shuqiang Zhang, Shuqiang Zhang, Sinong Wang, Sneha Agarwal, Soji Sajuyigbe, Soumith Chintala, Stephanie Max, Stephen Chen, Steve Kehoe, Steve Satterfield, Sudarshan Govindaprasad, Sumit Gupta, Summer Deng, Sungmin Cho, Sunny Virk, Suraj Subramanian, Sy Choudhury, Sydney Goldman, Tal Remez, Tamar Glaser, Tamara Best, Thilo Koehler, Thomas Robinson, Tianhe Li, Tianjun Zhang, Tim Matthews, Timothy Chou, Tzook Shaked, Varun Vontimitta, Victoria Ajayi, Victoria Montanez, Vijai Mohan, Vinay Satish Kumar, Vishal Mangla, Vlad Ionescu, Vlad Poenaru, Vlad Tiberiu Mihailescu, Vladimir Ivanov, Wei Li, Wenchen Wang, Wenwen Jiang, Wes Bouaziz, Will Constable, Xiaocheng Tang, Xiaojian Wu, Xiaolan Wang, Xilun Wu, Xinbo Gao, Yaniv Kleinman, Yanjun Chen, Ye Hu, Ye Jia, Ye Qi, Yenda Li, Yilin Zhang, Ying Zhang, Yossi Adi, Youngjin Nam, Yu, Wang, Yu Zhao, Yuchen Hao, Yundi Qian, Yunlu Li, Yuzi He, Zach Rait, Zachary DeVito, Zef Rosnbrick, Zhaoduo Wen, Zhenyu Yang, Zhiwei Zhao, and Zhiyu Ma. 2024. The Llama 3 Herd of Models. (2024). arXiv:cs.AI/2407.21783 https://arxiv.org/abs/2407.21783

[9] Fei Gui, Kaihui Gao, Li Chen, Dan Li, Vincent Liu, Ran Zhang, Hongbing Yang, and Dian Xiong. 2025. Accelerating design space exploration for LLM training systems with multi-experiment parallel simulation. In Proceedings of the 22nd USENIX Symposium on Networked Systems Design and Implementation (NSDI ’25). USENIX Association, USA, Article 25, 16 pages. [10] Hanpeng Hu, Chenyu Jiang, Yuchen Zhong, Yanghua Peng, Chuan Wu, Yibo Zhu, Haibin Lin, and Chuanxiong Guo. 2022. dPRO: A Generic Performance Diagnosis and Optimization Toolkit for Expediting Distributed DNN Training. In Proceedings of Machine Learning and Systems, D. Marculescu, Y. Chi, and C. Wu (Eds.), Vol. 4. 623–637. https://proceedings.mlsys.org/paper_files/paper/2022/file/ b422680f3db0986ddd7f8f126baaf0fa-Paper.pdf [11] Zhiyi Hu, Siyuan Shen, Tommaso Bonato, Sylvain Jeaugey, Cedell Alexander, Eric Spada, James Dinan, Jeff Hammond, and Torsten Hoefler. 2026. Demystifying NCCL: An In-depth Analysis of GPU Communication Protocols and Algorithms. (2026). arXiv:cs.DC/2507.04786 https://arxiv.org/abs/2507.04786 [12] Jinkun Lin, Ziheng Jiang, Zuquan Song, Sida Zhao, Menghan Yu, Zhanghan Wang, Chenyuan Wang, Zuocheng Shi, Xiang Shi, Wei Jia, Zherui Liu, Shuguang Wang, Haibin Lin, Xin Liu, Aurojit Panda, and Jinyang Li. 2025. Understanding stragglers in large model training using what-if analysis. In Proceedings of the 19th USENIX Conference on Operating Systems Design and Implementation (OSDI ’25). USENIX Association, USA, Article 27, 16 pages. [13] Guandong Lu, Runzhe Chen, Yakai Wang, Yangjie Zhou, Rui Zhang, Zheng Hu, Yanming Miao, Zhifang Cai, Li Li, Jingwen Leng, and Minyi Guo. 2023. DistSim: A performance model of large-scale hybrid distributed DNN training. In Proceedings of the 20th ACM International Conference on Computing Frontiers (CF ’23). Association for Computing Machinery, New York, NY, USA, 112–122. https://doi.org/10.1145/ 3587135.3592200 [14] Qingkai Meng, Hao Zheng, Zhenhui Zhang, ChonLam Lao, Chengyuan Huang, Baojia Li, Ziyuan Zhu, Hao Lu, Weizhen Dang, Zitong Lin, Weifeng Zhang, Lingfeng Liu, Yuanyuan Gong, Chunzhi He, Xiaoyuan Hu, Yinben Xia, Xiang Li, Zekun He, Yachen Wang, Xianneng Zou, Kun Yang, Gianni Antichi, Guihai Chen, and Chen Tian. 2025. Astral: A Datacenter Infrastructure for Large Language Model Training at Scale. In Proceedings of the ACM SIGCOMM 2025 Conference (SIGCOMM ’25). Association for Computing Machinery, New York, NY, USA, 609–625. https://doi.org/10.1145/3718958.3750521 [15] Meta. 2025. Holistic Trace Analysis. (2025). https://github.com/ facebookresearch/HolisticTraceAnalysis GitHub repository, latest release May 28, 2025. [16] NVIDIA Corporation. 2025. Github - CUDA Checkpoint and Restore Utility. https://github.com/NVIDIA/cuda-checkpoint. (2025). [17] OpenAI, Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, Red Avila, Igor Babuschkin, Suchir Balaji, Valerie Balcom, Paul Baltescu, Haiming Bao, Mohammad Bavarian, Jeff Belgum, Irwan Bello, Jake Berdine, Gabriel Bernadett-Shapiro, Christopher Berner, Lenny Bogdonoff, Oleg Boiko, Madelaine Boyd, Anna-Luisa Brakman, Greg Brockman, Tim Brooks, Miles Brundage, Kevin Button, Trevor Cai, Rosie Campbell, Andrew Cann, Brittany Carey, Chelsea Carlson, Rory Carmichael, Brooke Chan, Che Chang, Fotis Chantzis, Derek Chen, Sully Chen, Ruby Chen, Jason Chen, Mark Chen, Ben Chess, Chester Cho, Casey Chu, Hyung Won Chung, Dave Cummings, Jeremiah Currier, Yunxing Dai, Cory Decareaux, Thomas Degry, Noah Deutsch, Damien Deville, Arka Dhar, David Dohan, Steve Dowling, Sheila Dunning, Adrien Ecoffet, Atty Eleti, Tyna Eloundou, David Farhi, Liam Fedus, Niko Felix, Simón Posada Fishman, Juston Forte, Isabella Fulford, Leo Gao, Elie Georges, Christian Gibson, Vik Goel, Tarun Gogineni,

15

Gabriel Goh, Rapha Gontijo-Lopes, Jonathan Gordon, Morgan Grafstein, Scott Gray, Ryan Greene, Joshua Gross, Shixiang Shane Gu, Yufei Guo, Chris Hallacy, Jesse Han, Jeff Harris, Yuchen He, Mike Heaton, Johannes Heidecke, Chris Hesse, Alan Hickey, Wade Hickey, Peter Hoeschele, Brandon Houghton, Kenny Hsu, Shengli Hu, Xin Hu, Joost Huizinga, Shantanu Jain, Shawn Jain, Joanne Jang, Angela Jiang, Roger Jiang, Haozhun Jin, Denny Jin, Shino Jomoto, Billie Jonn, Heewoo Jun, Tomer Kaftan, Łukasz Kaiser, Ali Kamali, Ingmar Kanitscheider, Nitish Shirish Keskar, Tabarak Khan, Logan Kilpatrick, Jong Wook Kim, Christina Kim, Yongjik Kim, Jan Hendrik Kirchner, Jamie Kiros, Matt Knight, Daniel Kokotajlo, Łukasz Kondraciuk, Andrew Kondrich, Aris Konstantinidis, Kyle Kosic, Gretchen Krueger, Vishal Kuo, Michael Lampe, Ikai Lan, Teddy Lee, Jan Leike, Jade Leung, Daniel Levy, Chak Ming Li, Rachel Lim, Molly Lin, Stephanie Lin, Mateusz Litwin, Theresa Lopez, Ryan Lowe, Patricia Lue, Anna Makanju, Kim Malfacini, Sam Manning, Todor Markov, Yaniv Markovski, Bianca Martin, Katie Mayer, Andrew Mayne, Bob McGrew, Scott Mayer McKinney, Christine McLeavey, Paul McMillan, Jake McNeil, David Medina, Aalok Mehta, Jacob Menick, Luke Metz, Andrey Mishchenko, Pamela Mishkin, Vinnie Monaco, Evan Morikawa, Daniel Mossing, Tong Mu, Mira Murati, Oleg Murk, David Mély, Ashvin Nair, Reiichiro Nakano, Rajeev Nayak, Arvind Neelakantan, Richard Ngo, Hyeonwoo Noh, Long Ouyang, Cullen O’Keefe, Jakub Pachocki, Alex Paino, Joe Palermo, Ashley Pantuliano, Giambattista Parascandolo, Joel Parish, Emy Parparita, Alex Passos, Mikhail Pavlov, Andrew Peng, Adam Perelman, Filipe de Avila Belbute Peres, Michael Petrov, Henrique Ponde de Oliveira Pinto, Michael, Pokorny, Michelle Pokrass, Vitchyr H. Pong, Tolly Powell, Alethea Power, Boris Power, Elizabeth Proehl, Raul Puri, Alec Radford, Jack Rae, Aditya Ramesh, Cameron Raymond, Francis Real, Kendra Rimbach, Carl Ross, Bob Rotsted, Henri Roussez, Nick Ryder, Mario Saltarelli, Ted Sanders, Shibani Santurkar, Girish Sastry, Heather Schmidt, David Schnurr, John Schulman, Daniel Selsam, Kyla Sheppard, Toki Sherbakov, Jessica Shieh, Sarah Shoker, Pranav Shyam, Szymon Sidor, Eric Sigler, Maddie Simens, Jordan Sitkin, Katarina Slama, Ian Sohl, Benjamin Sokolowsky, Yang Song, Natalie Staudacher, Felipe Petroski Such, Natalie Summers, Ilya Sutskever, Jie Tang, Nikolas Tezak, Madeleine B. Thompson, Phil Tillet, Amin Tootoonchian, Elizabeth Tseng, Preston Tuggle, Nick Turley, Jerry Tworek, Juan Felipe Cerón Uribe, Andrea Vallone, Arun Vijayvergiya, Chelsea Voss, Carroll Wainwright, Justin Jay Wang, Alvin Wang, Ben Wang, Jonathan Ward, Jason Wei, CJ Weinmann, Akila Welihinda, Peter Welinder, Jiayi Weng, Lilian Weng, Matt Wiethoff, Dave Willner, Clemens Winter, Samuel Wolrich, Hannah Wong, Lauren Workman, Sherwin Wu, Jeff Wu, Michael Wu, Kai Xiao, Tao Xu, Sarah Yoo, Kevin Yu, Qiming Yuan, Wojciech Zaremba, Rowan Zellers, Chong Zhang, Marvin Zhang, Shengjia Zhao, Tianhao Zheng, Juntang Zhuang, William Zhuk, and Barret Zoph. 2024. GPT-4 Technical Report. (2024). arXiv:cs.CL/2303.08774 https://arxiv.org/abs/2303.08774 [18] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. 2019. PyTorch: An Imperative Style, HighPerformance Deep Learning Library. In Advances in Neural Information Processing Systems, H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (Eds.), Vol. 32. Curran Associates, Inc. https://proceedings.neurips.cc/paper_files/paper/2019/file/ bdbca288fee7f92f2bfa9f7012727740-Paper.pdf [19] Jianxing Qin, Jingrong Chen, Xinhao Kong, Yongji Wu, Tianjun Yuan, Liang Luo, Zhaodong Wang, Ying Zhang, Tingjun Chen, Alvin R. Lebeck, and Danyang Zhuo. 2026. Phantora: Maximizing Code Reuse in Simulation-based Machine Learning System Performance Estimation. In NSDI ’26.

[20] Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. 2020. ZeRO: memory optimizations toward training trillion parameter models. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis (SC ’20). IEEE Press, Article 20, 16 pages. [21] Saeed Rashidi, Srinivas Sridharan, Sudarshan Srinivasan, and Tushar Krishna. 2020. ASTRA-SIM: Enabling SW/HW Co-Design Exploration for Distributed DL Training Platforms. In 2020 IEEE International Symposium on Performance Analysis of Systems and Software (ISPASS). 81– 92. https://doi.org/10.1109/ISPASS48437.2020.00018 [22] Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He. 2020. DeepSpeed: System Optimizations Enable Training Deep Learning Models with Over 100 Billion Parameters. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining (KDD ’20). Association for Computing Machinery, New York, NY, USA, 3505–3506. https://doi.org/10.1145/3394486.3406703 [23] Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao. 2024. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. (2024). arXiv:cs.LG/2407.08608 https://arxiv.org/abs/2407.08608 [24] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. 2019. Megatron-lm: Training multibillion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053 (2019). [25] Srinivas Sridharan, Taekyung Heo, Louis Feng, Zhaodong Wang, Matt Bergeron, Wenyin Fu, Shengbao Zheng, Brian Coutinho, Saeed Rashidi, Changhai Man, et al. 2023. Chakra: Advancing performance benchmarking and co-design using standardized execution traces. arXiv preprint arXiv:2305.14516 (2023). [26] Borui Wan, Gaohong Liu, Zuquan Song, Jun Wang, Yun Zhang, Guangming Sheng, Shuguang Wang, Houmin Wei, Chenyuan Wang, Weiqiang Lou, Xi Yang, Mofan Zhang, Kaihua Jiang, Cheng Ren, Xiaoyun Zhi, Menghan Yu, Zhe Nan, Zhuolin Zheng, Baoquan Zhong, Qinlong Wang, Huan Yu, Jinxin Chi, Wang Zhang, Yuhan Li, Zixian Du, Sida Zhao, Yongqiang Zhang, Jingzhe Tang, Zherui Liu, Chuan Wu, Yanghua Peng, Haibin Lin, Wencong Xiao, Xin Liu, and Liang Xiang. 2025. Robust LLM Training Infrastructure at ByteDance. In Proceedings of the ACM SIGOPS 31st Symposium on Operating Systems Principles (SOSP ’25). Association for Computing Machinery, New York, NY, USA, 186–203. https://doi.org/10.1145/3731569.3764838 [27] Xizheng Wang, Qingxu Li, Yichi Xu, Gang Lu, Dan Li, Li Chen, Heyang Zhou, Linkang Zheng, Sen Zhang, Yikai Zhu, Yang Liu, Pengcheng Zhang, Kun Qian, Kunling He, Jiaqi Gao, Ennan Zhai, Dennis Cai, and Binzhang Fu. 2025. SimAI: Unifying Architecture Design and Performance Tuning for Large-Scale Large Language Model Training with Scalability and Precision. In 22nd USENIX Symposium on Networked Systems Design and Implementation (NSDI 25). USENIX Association, Philadelphia, PA, 541–558. https://www.usenix.org/conference/ nsdi25/presentation/wang-xizheng-simai [28] William Won, Taekyung Heo, Saeed Rashidi, Srinivas Sridharan, Sudarshan Srinivasan, and Tushar Krishna. 2023. ASTRA-sim2.0: Modeling Hierarchical Networks and Disaggregated Systems for Large-model Training at Scale. In 2023 IEEE International Symposium on Performance Analysis of Systems and Software (ISPASS). 283–294. https: //doi.org/10.1109/ISPASS57527.2023.00035 [29] An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, Chujie Zheng, Dayiheng Liu, Fan Zhou, Fei Huang, Feng Hu, Hao Ge, Haoran Wei, Huan Lin, Jialong Tang, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Yang, Jiaxi Yang, Jing Zhou, Jingren Zhou, Junyang Lin, Kai Dang, Keqin Bao, Kexin Yang, Le Yu, Lianghao Deng, Mei Li, Mingfeng Xue, Mingze Li, Pei Zhang, Peng Wang, Qin Zhu, Rui Men, Ruize Gao, Shixuan Liu, Shuang Luo, Tianhao Li, Tianyi Tang, Wenbiao Yin, Xingzhang Ren, Xinyu Wang, Xinyu Zhang, Xuancheng Ren, Yang 16

Fan, Yang Su, Yichang Zhang, Yinger Zhang, Yu Wan, Yuqiong Liu, Zekun Wang, Zeyu Cui, Zhenru Zhang, Zhipeng Zhou, and Zihan Qiu. 2025. Qwen3 Technical Report. (2025). arXiv:cs.CL/2505.09388 https://arxiv.org/abs/2505.09388 [30] Ted Zadouri, Markus Hoehnerbach, Jay Shah, Timmy Liu, Vijay Thakkar, and Tri Dao. 2026. FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling. (2026). arXiv:cs.CL/2603.05451 https://arxiv.org/abs/2603.05451 [31] Shulai Zhang, Ningxin Zheng, Haibin Lin, Ziheng Jiang, Wenlei Bao, Chengquan Jiang, Qi Hou, Weihao Cui, Size Zheng, Li-Wen Chang, Quan Chen, and Xin Liu. 2025. Comet: Fine-grained Computation-communication Overlapping for Mixture-of-Experts. (2025). arXiv:cs.DC/2502.19811 [32] Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, Alban Desmaison, Can Balioglu, Pritam Damania, Bernard Nguyen, Geeta Chauhan, Yuchen Hao, Ajit Mathews, and Shen Li. 2023. PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel. (2023). arXiv:cs.DC/2304.11277 https://arxiv.org/abs/2304.11277 [33] Hongyu Zhu, Amar Phanishayee, and Gennady Pekhimenko. 2020. Daydream: Accurately Estimating the Efficacy of Optimizations for DNN Training. In 2020 USENIX Annual Technical Conference (USENIX ATC 20). USENIX Association, 337–352. https://www.usenix.org/ conference/atc20/presentation/zhu-hongyu

17

A

COORDINATOR AND PRIORITY-BASED CONTEXT SWITCHING ALGORITHM.

Algorithm 1 Priority-Based GPU Context Switching 1: Input: Trigger rank 𝑟 𝑡 , Worker status 𝑆 2: Output: Best switching candidate 3: function SelectSwitch(𝑟 𝑡 ) 4: 𝑔𝑝𝑢 ← 𝑆 [𝑟𝑡 ].𝑔𝑝𝑢_𝑖𝑑

greedy scheduling strategy that maximizes system-wide progress by prioritizing processes with the highest potential for unblocking other waiting operations. Algorithm 1 presents the core switching logic. When a process encounters a collective operation that cannot be executed immediately, the coordinator must decide which frozen process to activate on the same GPU. The selection process (lines 4-19) applies several filtering criteria to identify viable candidates. First, it excludes finished processes and those pinned to different GPUs (lines 6-11), ensuring only relevant processes compete for the target GPU. Most critically, it filters out processes whose head-of-line blocking operation is not ready (lines 12-14), as switching to such processes would yield no forward progress. Among eligible candidates, the coordinator selects the process with the maximum number of pending operations (lines 16-19). This greedy heuristic is based on the insight that processes with more pending work are likely to unblock a larger number of dependent operations when given GPU time, thereby maximizing overall system throughput. The main coordination loop (lines 22-35) handles incoming collective notifications by first updating the pending operation counts for all waiting ranks (line 24). If all participants in a collective operation are currently resident on GPUs, the operation executes directly via the network proxy without context switching (lines 26-27). Otherwise, the coordinator performs context switching by freezing the requesting process, activating the selected candidate, and storing communication tensors to disk (lines 29-33) for later retrieval when all participants become available. This approach effectively transforms the limited GPU resources into a larger virtual execution environment, enabling trace generation for large-scale training scenarios using significantly fewer physical GPUs while maintaining correctness through careful dependency tracking and progressoriented scheduling.

B

𝑚𝑎𝑥_𝑝𝑒𝑛𝑑𝑖𝑛𝑔 ← 0, 𝑏𝑒𝑠𝑡 ← NULL for 𝑟𝑎𝑛𝑘 ∈ 𝑆 do if 𝑆 [𝑟𝑎𝑛𝑘].𝑠𝑡𝑎𝑡𝑢𝑠 = FINISHED then continue end if if 𝑆 [𝑟𝑎𝑛𝑘].𝑔𝑝𝑢_𝑖𝑑 ≠ 𝑔𝑝𝑢 then continue end if if 𝑆 [𝑟𝑎𝑛𝑘].ℎ𝑒𝑎𝑑_𝑜𝑝.𝑠𝑡𝑎𝑡𝑢𝑠 ≠ READY then con11: tinue 12: end if 13: if 𝑆 [𝑟𝑎𝑛𝑘].𝑝𝑒𝑛𝑑𝑖𝑛𝑔_𝑜𝑝𝑠 > 𝑚𝑎𝑥_𝑝𝑒𝑛𝑑𝑖𝑛𝑔 then 14: 𝑚𝑎𝑥_𝑝𝑒𝑛𝑑𝑖𝑛𝑔 ← 𝑆 [𝑟𝑎𝑛𝑘].𝑝𝑒𝑛𝑑𝑖𝑛𝑔_𝑜𝑝𝑠 15: 𝑏𝑒𝑠𝑡 ← 𝑟𝑎𝑛𝑘 16: end if 17: end for 18: return 𝑏𝑒𝑠𝑡 19: end function 5: 6: 7: 8: 9: 10:

20: 21: function HandleCollective(𝑚𝑠𝑔)

𝑟 ← 𝑚𝑠𝑔.𝑠𝑒𝑛𝑑𝑒𝑟 UpdatePendingOps(𝑚𝑠𝑔.𝑟𝑒𝑚𝑎𝑖𝑛𝑖𝑛𝑔_𝑟𝑎𝑛𝑘𝑠) if AllParticipantsOnGPU(𝑚𝑠𝑔.𝑝𝑎𝑟𝑡𝑖𝑐𝑖𝑝𝑎𝑛𝑡𝑠) then 25: ExecuteDirect(𝑚𝑠𝑔) 26: else 27: 𝑐𝑎𝑛𝑑𝑖𝑑𝑎𝑡𝑒 ← SelectSwitch(𝑟 ) 28: if 𝑐𝑎𝑛𝑑𝑖𝑑𝑎𝑡𝑒 ≠ NULL then 29: FreezeRank(𝑟 ) 30: ActivateRank(𝑐𝑎𝑛𝑑𝑖𝑑𝑎𝑡𝑒) 31: StoreTensors(𝑚𝑠𝑔.𝑡𝑒𝑛𝑠𝑜𝑟𝑠) 32: end if 33: end if 34: end function 22: 23: 24:

35: 36: function UpdatePendingOps(𝑤𝑎𝑖𝑡𝑖𝑛𝑔_𝑟𝑎𝑛𝑘𝑠) 37: for 𝑟𝑎𝑛𝑘 ∈ 𝑤𝑎𝑖𝑡𝑖𝑛𝑔_𝑟𝑎𝑛𝑘𝑠 do 38: 𝑆 [𝑟𝑎𝑛𝑘].𝑝𝑒𝑛𝑑𝑖𝑛𝑔_𝑜𝑝𝑠

CONTEXT SWITCHING OVERHEAD

𝑆 [𝑟𝑎𝑛𝑘].𝑝𝑒𝑛𝑑𝑖𝑛𝑔_𝑜𝑝𝑠 + 1 39: end for 40: end function

While the context-switching-based approach enables graph collection with limited GPUs, it introduces non-trivial overhead. First, storage overhead can be substantial. For example, an A100 GPU has 80GB of memory; replaying a 128GPU job would require storing up to 128 × 80GB (≈ 10 TB) of data in CPU memory for swapping, leading to significant memory pressure. To scale beyond CPU memory limits, PrismLLM can spill inactive contexts to disk using CRIU [3]. Second, GPU context switching is expensive. For an 80GB GPU, checkpointing and restoring a CUDA context can take more than two minutes using existing tools (e.g., cuda-checkpoint [16]). 18

←

C

Table 2. Model structures evaluated.

EXAMPLE OF COMMUNICATION INPUT TENSORS GENERATE RULES

PrismLLM allows advanced users with full control and understanding of their program directly specify the tensors required for communication, thereby avoiding the overhead of context-switching-based execution. We present the generate rules we used in our experiments as an example. We designed these rules because the data within these specific tensors affects the program’s execution logic and control flow, rather than just participating in numerical computations.

Model

Params

Layers

Heads

Experts (TopK)

M1 M2 M3

235BA22B 503BA20B 1.01TA43B

94 62 62

64 32 64

128 (8) 256 (8) 256 (8)

Table 3. Parallelization strategies evaluated.

C.1 Dataloader statuses During Megatron-LM framework initialization, dataloaders are constructed across multiple ranks. Following construction, a broadcast is performed to synchronize the status of the rank 0 dataloader. If rank 0 fails to initialize, the training pipeline terminates. In our experiments, we inject data representing a "successful" status into this communication to ensure the emulation proceeds through all training steps.

Strategy

TP

PP

VPP

EP

GA

A B C D

1 2 1 1

4 4 16 8

0 2 0 0

8 8 8 16

8 16 32 16

*TP: Tensor Parallel, PP: Pipeline Parallel, VPP: Virtual Pipeline Parallel (number of virtual stages per pipeline stage; 0 = disabled), EP: Expert Parallel, GA: Gradient Accumulation.

Sandbox as Root: In the reduction (upward) stage, the sandbox must compute the final fully reduced value for the chunk. Therefore, its child vRank is instructed to prepare and send a compensated value datafull − datasandbox , where datasandbox represents the aggregated contribution of all ranks within the sandbox. In the subsequent broadcast (downward) stage, the sandbox already holds the correct result and simply propagates it downward; the child vRanks passively receive it without affecting the sandbox’s correctness. Sandbox as Leaf: In the reduction stage, the sandbox’s proxy rank simply forwards its locally aggregated data upward to its parent vRank. Since the final reduction occurs elsewhere, the specific value is irrelevant to the sandbox at this point. In the broadcast stage, however, the sandbox must receive the final result. Thus, its parent vRank actively sends the correct datafull downward to the sandbox. Sandbox as Intermediate: For an intermediate sandbox node, its locally reduced value will ultimately be overwritten by the broadcast phase. During reduction, its child vRank can send arbitrary values (ANY). During broadcast, its parent vRank is responsible for sending datafull downward, ensuring the sandbox and its intra-node ranks receive the correct final tensor. By applying these localized rules at the sandbox boundary, our approach guarantees numerical correctness for treebased all-reduce while maximizing the use of arbitrary values (ANY) to minimize the overhead of mocking data.

C.2 Training samples When tensor model parallel (TP) is enabled in Megatron-LM, ranks that host the model’s embedding layer within the same TP group receive input data via a broadcast from rank 0. This data is used as indices for embedding lookups. If the indices exceed the vocabulary size, an index out of bounds occurs. In our experiments, we inject valid, in-vocab values into this communication. C.3 MoE Dispatch All-to-All Splits In Megatron Core’s MoE implementation, an allgather operation is performed before the dispatch phase to collect gating results from all ranks. These results are used to calculate the split sizes for subsequent all-to-all communication and to pre-allocate communication buffers. Overly large split values can lead to unintended OOM errors. We address this by injecting data that simulates "zero-data" transmissions from peer ranks, ensuring the calculated splits remain within a reasonable range and preventing OOM.

D EXTENSION TO TREE-BASED ALGORITHMS We now extend our approach to tree-based collective algorithms, such as the Double-Tree topology used in NCCL. Unlike the ring algorithm, NCCL constructs hierarchical trees for inter-node communication, while intra-node communication typically forms a chain. In the tree topology, inter-node connections are established between specific proxy ranks. Therefore, we ensure correctness by manipulating the data exclusively at the boundary where a sandbox proxy rank connects to an external virtual rank (vRank). For a given data chunk, the tree topology assigns the sandbox proxy rank one of three roles: Root, Leaf, or Intermediate. Based on this role, the neighboring vRanks adjust their communication as follows:

E

MODEL PARAMETERS AND PARALLELIZATION STRATEGIES

Our experiments utilize an open-source, Megatron-LM-based implementation of the Qwen 3 MoE pretraining framework. The detailed model architectures are summarized in Table 2. For each model configuration, we evaluate two distinct parallelization strategies, with specific parameters detailed in Table 3. Furthermore, Megatron’s distributed optimizer is enabled across all experimental setups. 19

F

MOE MOCK ROUTER

Section 8.2, PrismLLM gracefully resolves this hardware bottleneck by scaling out virtual ranks across multiple assistant nodes to restore emulation fidelity.

For MoE models, memory allocation typically varies across devices because experts on different devices may process uneven amounts of data after the dispatch phase. The final data distribution is determined by the gating mechanism. After the input data on each rank passes through the gating layer, each rank obtains the probabilities (or logits) of its local tokens being dispatched to each expert. All ranks participating in expert model parallelism then perform communications to exchange these local probabilities. By combining these probabilities with the routing strategy (e.g., top_k), the final global data distribution is determined. To control the non-uniform dispatching of MoE models and observe the resulting memory footprint after dispatch, we designed the MoE Mock Router. This component utilizes the Balance Ratio (br) to regulate distribution probabilities. The br represents the ratio of the actual data volume possessed by a specific rank to the volume it would possess under a perfectly uniform distribution. A br > 1 indicates that the rank is over-utilized relative to the average. Since each microbatch must pass through all layers within a single iteration, multiple gating operations occur, each requiring control via br. We characterize the distribution of br in one iteration using statistical metrics including br_min (minimum), br_max (maximum), br_avg (average), br_std (standard deviation), br_med (median), and br_skew (skewness). Based on these statistics, the MoE Mock Router derives the br distribution and pre-calculates the logits for all ranks within an expert model parallel group. These pre-calculated logits are subsequently injected into the logits tensor generated by the gating mechanism during each invocation. To ensure that these pre-calculated values do not incur additional GPU memory overhead, they are pinned in host memory. The Mock Router performs an asynchronous inplace copy to the GPU-resident logits tensor only when the gating operation is triggered, effectively overwriting the original values without allocating new device buffers.

G

H

RESULTS OF PRISMLLM VS. SIMAI

Figure 15 presents the comparison between PrismLLM and SimAI across all model and strategy configurations used in Section 8.4. SimAI is excluded from the memory evaluation due to its lack of memory allocation simulation support. SimAI fails to achieve acceptable fidelity, exhibiting a significant average underestimation error of 77.2%.

LATENCY EVALUATION OF ALLREDUCE

Table 4 presents an extended evaluation of the AllReduce transmission latency, detailing performance across varying message sizes (from 16 MB to 32 GB) and expanding the emulation scale up to 128 ranks. The results highlight the severe resource contention inherent in the Vanilla emulation approach, which suffers from exponential latency inflation as the message size and cluster scale grow. In contrast, PrismLLM’s communication pruning strategy consistently minimizes overhead, closely tracking the physical baseline across all tested configurations. However, as the scale reaches 128 ranks, PrismLLM exhibits noticeable latency deviations for large messages (e.g., 267.42 ms vs. 177.62 ms at 32 GB). This is expected: hosting all virtual ranks on a single physical node saturates its PCIe bandwidth and SMs. As detailed in 20

Table 4. Extended evaluation of AllReduce transmission latency (ms) across different message sizes and cluster scales.

Baseline

Size

PrismLLM

32

64

128

16

32

64

128

16

32

64

128

0.27 1.22 6.32 47.72 189.09

0.58 1.17 5.65 43.61 173.25

0.41 1.56 6.40 44.34 177.08

0.50 2.17 8.11 45.04 177.62

0.28 1.22 6.38 47.92 189.22

85.92 92.98 208.17 1663.51 6655.51

181.98 734.74 1070.08 4256.79 16950.36

639.44 981.97 4351.50 12742.20 50815.80

0.27 1.18 6.19 47.21 187.96

0.55 1.12 5.56 43.34 173.11

0.27 1.82 6.34 44.71 176.70

0.97 3.97 12.72 67.04 267.42

Baseline

PrismLLM

SimAI

-79.3 0.37% 7% + -77.6 0.81% 1% + -78.4 1.54% 5% -75.8 0.24% 1% -71.7 0.02% 4% +0.1 7% -79.8 6%

Figure 15. End-to-end iteration time estimation results compared with SimAI.

21

+0.6 0 8% % -78.8

-78.6

+0.2 6% 1%

Scale: 2048 GPUs 12 12 12 9 9 9 6 6 6 3 3 3 0 235B 235B 503B 503B 1.01T 1.01T 0 235B 235B 503B 503B 1.01T 1.01T 0 235B 235B 503B 503B 1.01T 1.01T (S.A) (S.B) (S.A) (S.B) (S.C) (S.D) (S.A) (S.B) (S.A) (S.B) (S.C) (S.D) (S.A) (S.B) (S.A) (S.B) (S.C) (S.D) +0.8 5 3% %

Scale: 1024 GPUs

-79.9

Scale: 512 GPUs

-75.5 1.98% 0% + -72.1 0.48% 1% -0.40 -79.7 % 4%

16

-79.3 0.62% 9% -78.3 0.52% 8% -78.9 0.12% 4% -74.5 0.92% 8% -71.6 0.15% 0% +0.4 8% -79.3 5%

Iteration Time (s)

16M 128M 1G 8G 32G

Vanilla Emulation