arXiv CSopen access

Simulating Unified Tensor Resharding in heterogeneous AI systems

Simulating Unified Tensor Resharding in heterogeneous AI systems Sumit Kumar1 , Sayantan Dasgupta1 , Kushal Mitra1 , Meet Dadhania2 , Rohan Sudhir Basugade1 , Praveen Tammana2 , Satananda Burla3 , Abed Mohammad Kamaluddin4 , Rinku Shah1 1 IIIT-Delhi, India

2 IIT Hyderabad, India

3 Marvell Technology Inc., USA

arXiv:2606.26633v1 [cs.DC] 25 Jun 2026

Abstract

Table 1. AI infrastructure at hyperscalers.

State-of-the-art AI training simulators assume homogeneous compute and network infrastructure. However, real-world training infrastructure is becoming increasingly heterogeneous since: (a) Model architectures such as multimodal and MoE exploit heterogeneity to improve device utilization, (b) Public cloud platforms often provide limited availability of homogeneous hardware due to fast hardware evolution, and (c) Large enterprises frequently deploy geographically distributed infrastructure that is both diverse and heterogeneous. In this paper, we present 𝑋𝑠𝑖𝑚, a heterogeneity-aware simulator for distributed LLM training. 𝑋𝑠𝑖𝑚 supports: (i) Load balancing through non-uniform workload partitioning across heterogeneous device groups, (ii) Heterogeneity-aware collective communication via customized ring construction and chunk partitioning, (iii) Reusable heterogeneity-aware abstractions for emerging pipeline-parallel algorithms and non-uniform tensor resharding technique, (iv) Flexible input abstractions for specifying deployment plans with custom device groups and custom device-to-parallelism mappings, and (v) Pluggable integration with NS-3 and htsim, allowing users to trade off simulation fidelity for performance and scalability. Our evaluation demonstrates that 𝑋𝑠𝑖𝑚 accurately predicts training time for real-world heterogeneous deployments, with an error of less than 5% across most heterogeneous data-parallel/tensor-parallel configurations and around 2% error with pipeline-parallel communication modeling. We expose actionable metrics such as pipeline bubble time and straggler waiting time.

1

4 Marvell Technology Inc., India

HyperscalerTopology

Fabric

Google Cloud [9]

Railoptimized

GPUDirect A4X Max (GB300 Ultra superchips), RDMA, GPUDi- A3 Ultra (NVIDIA H200 GPUs), A3 rect TCPX. High (NVIDIA H100 GPUs).

Compute and Interconnect

Alibaba Cloud [8]

SHENLONG Hybrid RDMA Heterogeneous nodes combining architecand VPC net- GPUs (L20, A10, V100, T4, P100), ture [7] working [49] FPGAs, and ASICs.

Meta [18]

Two-stage Clos

Microsoft Azure [23]

Full mesh, fat RDMA [3] tree

RoCEv2 [38]

ZionEX with NVIDIA A100 GPUs, Grand Teton NVIDIA H100 GPUs. GB300 NVL72 systems.

Table 2. Evolution of NVIDIA data center GPUs. Model

Year

Archit.

CUDA Cores

Memory (GB)

FP16 (TFLOPS)

A100 H100 H200 B100 B200

2020 2022 2023 2024 2024

Ampere Hopper Hopper Blackwell Blackwell

6912 16896 16896 14592 16896

40/80 80 141 192 192

77.97 204.9 989.5 1800 2250

Understanding distributed training behavior at scale necessitates a simulator. The large design space induces complex interactions across workload, system, and network layers, making it difficult to isolate the impact of individual design choices in production clusters. Moreover, highimpact events such as congestion, stragglers, and tail-latency amplification are difficult to reproduce consistently in real deployments [53]. These challenges motivate the need for a simulator that can systematically explore the design space and stress-test AI infrastructure under controlled, repeatable conditions. State-of-the-art (SOTA) simulators such as ASTRAsim [62], SimAI [60], Multiverse [21], ATLAHS [50] and Meta’s Arcadia [16], provide a full-stack simulation framework for training clusters. However, they assume homogeneous compute and networking infrastructure. This assumption does not reflect modern AI deployments (see Table 1). As a result, these simulators risk obscuring critical interactions, leading to misleading conclusions and suboptimal infrastructure designs [55, 65]. In this paper, we present 𝑋𝑠𝑖𝑚, a system-level LLM training simulator for real-world AI clusters with heterogeneous compute and network infrastructure. Modern distributed training deployments are increasingly becoming heterogeneous. Heterogeneity in distributed AI training infrastructure arises naturally from four

Introduction

Rapid evolution and scale in modern distributed training systems create a complex design space spanning parallelism strategies, transport protocols, network topologies, and hardware configurations [62]. To meet the workload demands at scale, leading hyperscalers tightly co-design network fabric, network topology, and accelerator interconnects for efficient distributed AI training. For example, Meta [18] deploys dedicated RoCEv2 [38] fabrics with two-stage Clos topology to provide non-blocking, low-latency collective communication [52] across GPU clusters. Microsoft Azure [23] combines high-bandwidth intra-rack NVLink with InfiniBand interrack connectivity to scale GPU clusters. 1

primary factors. (a) GPU compute and interconnect bandwidth advance approximately 3.0× and 1.4× annually [19], respectively (see Table 2). Organizations progressively integrate new GPU generations into existing clusters before retiring older hardware, resulting in a transitional heterogeneous state that may persist for months or years; (b) Demand for larger AI models continues to surge, for example, Meta’s Llama-4 with 2 trillion parameters was trained on 32𝐾 GPUs [1]. With power regulations within a given region, enterprises such as Meta run AI training across multiple data centers, each equipped with distinct hardware capabilities and interconnect architectures [34]; (c) Since the supply of GPU resources is unable to meet the demand, some organizations rely on the cloud resources. Cloud users often cannot wait for uniform GPU allocations; instead, they accept heterogeneous configurations to meet time-to-market constraints [26, 35, 48, 61]. Table 1 shows the configurations offered by popular hyperscalers; (d) Device heterogeneity is leveraged for better utilization. For instance, in multimodal model training, compute-intensive modalities (LLM backbone) are mapped to newer GPUs, and less computeintensive tasks (encoder) are mapped to older hardware [72]. Can we naively reuse a homogeneous deployment plan for a heterogeneous cluster? We observe three major roadblocks: (1) Equal-sized model or data partitions when mapped to devices with asymmetric compute and memory result in load imbalance, causing faster GPUs to wait for slower ones. Metis [58] reported an 8× slowdown in training throughput due to load imbalance. (2) Adopting unequal-sized model partitioning helps match the processing speeds. However, (a) Collective operations such as AllReduce must handle misaligned batch sizes and tensor shapes [44], and (b) Network traffic generated for gradient synchronization is deterministic but not uniform, resulting in the need for custom transport protocols to reduce training completion time; (3) Network bottlenecks intensify as gradient and activation communications traverse hybrid interconnects. Motivating Example: In our evaluation of the Llama-7B model [57], SimAI [60], which assumes homogeneous clusters, mispredicts training time for a heterogeneous deployment by up to 80%, whereas 𝑋𝑠𝑖𝑚 maintains an error below 4.7% (see Fig. 6). This gap arises because SimAI does not capture straggler effects caused by non-uniform compute capacities during collective communication. Resharding Under Arbitrary Heterogeneity. Recent work has proposed heterogeneous collective communication [25, 73] and specialized resharding mechanisms for layout mismatches [2, 44, 75]. However, these approaches target narrow scenarios and assume fixed communication structures, limiting their applicability to modern heterogeneous training deployments that combine non-uniform data-parallel (DP), tensor-parallel (TP), and pipeline-parallel (PP) configurations. 𝑋𝑠𝑖𝑚 bridges this gap with a unified

Table 3. Comparison of state-of-the-art distributed training Features Trace extrapolation support Full stack training simulation Collective optimization support Pipeline parallelism support (discrete-event) Network protocol simulation Heterogeneous cluster simulation

Astra- Echo vTrain MultiATLAHS SimAI 𝑋𝑠𝑖𝑚 Sim [62] [17] [5] verse [21] [50] [60] ✗

✓

✗

✗

✗

✓

✓

✓

✓

✓

✓

✓

✓

✓

✗

✓

✓

✗

✓

✓

✓

✗

✗

✗

✗

✓

✗

✓

✓

✗

✗

✓

✓

✓

✓

✗

✗

✗

✗

✗

✗

✓

abstraction enabling efficient synchronization across arbitrary heterogeneous parallelism strategies. Evaluating training performance of heterogeneous deployment plan. The core idea of stateof-the-art heterogeneity-aware deployment strategies [27, 39, 58, 64, 65, 71] is to non-uniformly partition the model across heterogeneous GPUs to improve resource utilization and performance [65]. To evaluate the proposed deployment plan for heterogeneous clusters, researchers rely on: (a) real-world deployments [27, 58, 65], which do not have optimized collective communication implementation, and are not scalable and accessible to all, or (b) analytical simulations [55, 56], which do not mimic real-world conditions, or end-to-end full system simulation. These observations motivate the need for a new simulator framework that incorporates heterogeneity-aware deployment strategies and accurately models both compute and network diversity for large-scale model deployment. Our proposal. We present 𝑋𝑠𝑖𝑚, a system-level simulator for heterogeneous AI training that models clusters with heterogeneous compute nodes1 and heterogeneous interconnects2 . 𝑋𝑠𝑖𝑚 introduces abstractions for non-uniform workload partitioning, heterogeneity-aware tensor resharding and synchronization, and pipeline-parallel execution with straggler effects. It further extends both NS-3 [4] and htsim [14] to model heterogeneous scale-up and scale-out networks, providing a choice between high-fidelity packetlevel protocol simulation and scalable flow-level simulation. Together, these capabilities enable end-to-end simulation of heterogeneous training deployments spanning compute, communication, and network layers. Key Contributions. To the best of our knowledge, 𝑋𝑠𝑖𝑚 is the first heterogeneity-aware, system-level simulator for AI training (see Table 3). Our key contributions include: (1) We identify core requirements for heterogeneous AI training simulation (§3) and design abstractions to specify nonuniform workload partitioning across heterogeneous device groups (§4.1). 1We assume homogeneous GPUs within a node, and heterogeneous GPUs

across nodes and clusters to mimic real-world infrastructure. 2We assume scale-up heterogeneity across nodes/clusters, whereas scale-out heterogeneity across clusters.

2

(2) We design Sweep-line-based DP grouping, and LCMbased multi-ring collectives and chunking algorithms to enable synchronization and resharding across mismatched DP, TP, and PP layouts (§4.3, §4.4). We design primitives for collective communication and tensor resharding, which we leveraged to implement two SOTA tensor resharding frameworks: HetAuto [44] and AlpaComm [75]. (3) We extend the runtime to model heterogeneous pipelineparallel execution (e.g., pipeline barriers, GPipe-style communication, straggler effects), providing reusable primitives for future PP strategies (§4.5). (4) We extend NS-3 and htsim network simulators to support heterogeneous scale-up/scale-out networks, capturing variations in link bandwidth, PCIe interconnects, NIC/scale-up delays, and topologies (§4.6). (5) We demonstrate that 𝑋𝑠𝑖𝑚 is accurate and scalable: it predicts training time within 7% error for NVLink, predicts training time within 17% error for hybrid (NVLink, PCIe, and Ethernet in non-isolated cluster), speeds up network simulation by up to 47× via htsim, and exposes actionable metrics (e.g., pipeline bubble time, straggler waiting time, TCO (Total Cost of Ownership)) for capacity planning (§5). 𝑋𝑠𝑖𝑚’s heterogeneity-aware resharding algorithm reduces synchronization overhead by up to 21% compared to the state of the art.

2

Background and Motivation

2.1

Emergent Challenges with Heterogeneity in AI Training Clusters

researchers to validate their workload distributions that maximize cluster utilization before deploying to real hardware. Challenge 2: Synchronization Overhead and Tensor Shape Mismatches During gradient synchronization, asymmetric workload partitioning leads to mismatches in tensor shape and batch size. This requires additional resharding operations to restore compatibility. Collective operations, such as AllReduce or broadcast, must combine/split varying tensor shapes and batch sizes across device groups with different degrees of tensor parallelism (TP) and Data parallelism (DP). Heterogeneity-aware training solutions repartition and redistribute tensor data, which incurs: (a) computation overhead: to split/combine gradient tensor before participating in collective operation, and (b) communication overhead: to exchange reshaped tensor segments using custom collective operations, i.e., multigroup DP ring 4 , which adds network traffic beyond standard gradient synchronization. How can a simulator help? A heterogeneity-aware simulator enables deployers to decide an optimal deployment plan after evaluating varying TP degrees, chunking strategies, and communication topologies (such as multi-ring configurations) without accessing vast, expensive physical clusters. Challenge 3: Communication Bottlenecks in Heterogeneous Networks Modern clusters deploy heterogeneous interconnects: NVLink (Gen3: 600 GB/s vs Gen4: 900 GB/s), PCIe (Gen4: 512 Gbps vs Gen5: 1024 Gbps), and RDMA NICs (ConnectX-6: 200 Gbps vs ConnectX-7: 400 Gbps ) [40]. In blocking collective communication operations such as AllReduce, networks with slower links significantly increase communication tail latency, leading to a substantial increase in end-to-end training performance [65]. How can a simulator help? A heterogeneity-aware simulator with network modeling can capture asymmetric bandwidth, latency, congestion, and tail latencies across mixed interconnects. By simulating communication patterns prior to deployment, deployers can identify bottlenecks in specific heterogeneous configurations, evaluate alternative topology mappings, and predict which interconnect placements or workload partitioning strategy can minimize communication overhead.

We observe heterogeneity-induced challenges in utilization, deployment complexity, and communication. Challenge 1: Load Balancing Across Asymmetric Hardware. Uniform workload partitioning is the default strategy in frameworks like Megatron-LM [51] and DeepSpeed [47]. Training frameworks assign equal portions of model layers, tensor slices, or data batches to all devices regardless of their compute capabilities. In heterogeneous clusters, this approach causes severe performance degradation due to compute imbalance. That is, high-performance GPUs (e.g., H100 with 204.9 FP16 TFLOPS) complete their assigned work significantly faster than lower-tier GPUs (e.g., A100 with 77.97 FP16 TFLOPS), leading to substantial idle time during synchronization barriers. Heterogeneity-aware training frameworks such as HexiScale [65] and Metis [58] use non-uniform partitioning strategies that assign layer/tensor/batch slices proportional to the device’s compute and memory capabilities, and create custom device groups (DG) 3 . How can a simulator help? A heterogeneity-aware simulator can track GPU idle time during synchronization barriers, and evaluate different partitioning strategies. This enables

2.2

Heterogeneity-aware AI Training Deployment: An Example

State-of-the-art heterogeneity-aware AI training systems jointly optimize device grouping, hybrid parallelism (PP, TP, DP), and non-uniform partitioning of data, layers, and tensors. By mapping each device group to an appropriate parallelism configuration and degree, these systems improve resource utilization and reduce end-to-end training time. 4 A multigroup DP ring is a data-parallel communication ring that spans

3 A device group represents a set of compute nodes with homogeneous

devices belonging to multiple device groups, enabling gradient synchronization across heterogeneous hardware groups.

compute and interconnect properties, modeled as a single unit for assigning parallelism strategies.

3

Figure 2. Comparison of SOTA tensor resharding approaches for transferring a 12-element global tensor from a source stage (𝑇 𝑃 = 6) to a destination stage (𝑇 𝑃 = 4). (a) HetAuto employs a 3-phase hierarchical approach (Gather → P2P → Scatter) that groups devices into two virtual clusters, each governed by GCD(6, 4) = 2, and routes traffic through leaders. (b) AlpaComm establishes direct point-to-point connections based on irregular cutpoint intervals, leading to non-uniform chunks.

Figure 1. Example deployment with heterogeneous compute. 7B model is partitioned across two nodes, i.e., one node with 5×H100 ("H") and the other with 5×A100 ("L"). The model and data batches are non-uniformly partitioned to efficiently utilize resources. Figure 1 illustrates a hybrid parallelism configuration for training the Llama-2 (7B) [57] model with 32 layers on a heterogeneous cluster. The cluster consists of two nodes: 𝑁𝑜𝑑𝑒𝐴 , equipped with 5 × 𝐻 100 80G (in blue), and 𝑁𝑜𝑑𝑒𝐵 , equipped with 5 × 𝐴100 40G GPUs (in red). Non-uniform workload partitioning. To effectively utilize this heterogeneous setup (Fig. 1), the workload is nonuniformly partitioned across device groups, with different batch sizes across DP groups (i.e., 8 and 16 are marked as green), different tensor-parallel degrees (e.g., TP = 2 and TP = 3), and asymmetric pipeline partitioning (i.e., 20 layers assigned to 𝐷𝐺 0 and 12 layers to 𝐷𝐺 1 marked as yellow). Resharding. During the backward pass, gradient synchronization across device groups must be preceded by resharding to ensure a common tensor shape and partitioning. Resharding is unavoidable when communicating groups use different TP degrees. Figure 1 shows that all inter-node communication between H100 and A100 device groups involves mismatched TP degrees and therefore requires resharding before collective operations. In contrast, pipeline parallelism (PP), if implemented in isolation, does not require resharding, as inter-stage communication is point-to-point and sequential despite differing group sizes. Some heterogeneity-aware systems (e.g., HetPipe [43]) avoid resharding by adopting a parameter-server–based synchronization model, at the cost of increased communication and potential server-side bottlenecks. 2.3

communication, and memory operations to generate a realworld workload trace. Further, SimAI [60], Echo [17], and Multiverse [21] extrapolate the real-world traces obtained for a subset of the GPU cluster, whereas NeuSight [29] and vTrain [5] are simulators that analytically predict training time. AstraSim [62], SimAI [60], ATLAHS [50], and Multiverse [21] simulate the entire training stack, including the workload generator, workload partitioning for parallelism, partition assignment to the set of GPUs (i.e., device groups), collective communication management, generation of compute and communication events, event scheduling, simulating compute events across GPUs, and simulating communication events over the training cluster. However, none of the existing works support LLM training in a heterogeneous cluster setting that comprises varied device (GPU and interconnect) capabilities and configurations. 𝑋𝑠𝑖𝑚 fills this gap. 2.4

State-of-the-art tensor resharding approaches

Heterogeneous LLM training often uses different DP, TP, and PP degrees across device groups, creating tensor layout mismatches that require resharding [44, 58, 65]. Prior work [44, 75] proposes optimized resharding mechanisms for these layouts. In this section, we provide an overview of these works using an example. Fig. 2 (a) shows HetAuto [44] addresses resharding across heterogeneous pipeline stages, i.e., 𝑇 𝑃 = 6 → 𝑇 𝑃 = 4, using a GCD-based three-phase strategy. It partitions source and destination devices into virtual groups based on the GCD of their TP/DP degrees. For example, between 𝑇 𝑃 = 6 and 𝑇 𝑃 = 4 stages, HetAuto uses GCD(6, 4) = 2 to form two virtual groups and coordinate data movement through leader nodes. Data is then reorganized through three phases: (i) intra-cluster gather to leader nodes, (ii) cross-cluster P2P

State-of-the-art AI Training Simulators

State-of-the-art simulators (see Table 3) can be classified into two broad categories: (1) workload generators [54] and (2) full-stack training simulators [17, 21, 50, 60, 62]. Chakra[54] executes real-world training workload on GPUs to capture the dependencies between compute, 4

transfer between leaders, and (iii) intra-cluster scatter to destination ranks. To resolve tensor partitioning mismatches, AlpaComm [75] applies a cutpoint-union algorithm that computes the union of shard boundaries from both layouts, partitioning the global tensor into atomic communication units. For example, when resharding a 12-element tensor from 𝑇 𝑃 = 6 to 𝑇 𝑃 = 4 (Fig. 2 (b)), the source boundaries {0, 2, 4, 6, 8, 10, 12} and destination boundaries {0, 3, 6, 9, 12} produce a union boundary set of {0, 2, 3, 4, 6, 8, 9, 10, 12}. These boundaries divide the tensor into eight irregular, non-uniform communication units [2, 1, 1, 2, 2, 1, 1, 2], which are then explicitly mapped between sender and receiver ranks.

3

optimal ring for gradient synchronization in a heterogeneous cluster. The algorithm must automatically identify valid DP groups (which may overlap) and select an optimal DP communication group to form the ring. [C3] Heterogeneity-aware Chunk Partitioning (§4.4). To handle tensor shape mismatches caused by heterogeneous TP sizes, the simulator must implement an algorithm that appropriately partitions the gradients into chunks, groups them, and shares them during collective operations, such as AllReduce. [C4] Pipeline-Parallel Communication and Synchronization. (§4.5) A heterogeneity-aware simulator must faithfully model pipeline-parallel dependencies, including pipeline barriers and inter-stage activation/gradient communication. It must also quantify straggler effects and pipeline bubbles arising from non-uniform compute capabilities, parallelism degrees, and communication paths across stages. [C5] Scalable Network Simulation (§4.6). To verify largescale heterogeneous deployments efficiently, the simulator should integrate network simulators that support flow-level network simulation. This approach enables modeling Flow Completion Time (FCT), congestion, and tail latencies for thousands of flows without the computational overhead of packet-level simulation. [C6] Heterogeneous Compute Simulation (§4.5). The core simulation engine must capture and model execution time based on device-specific capabilities. This involves capturing and simulating per-layer computation time scaled by the specific GPU type of the assigned device group.

Design Requirements

Heterogeneity-aware distributed AI training solutions such as Hexiscale [65], Metis [58], FlashFlex [66], Cephalo [22], and Atlahs [50] solve complex combinatorial optimization problems to generate a deployment for training within a heterogeneous cluster. The deployment comprises: (a) a heterogeneous compute cluster, (b) a heterogeneous cluster and host topology for scale-out, and scale-up, respectively, and (c) asymmetric workload to device mapping. Among other things, a heterogeneity-aware AI training simulator must have the provision to ingest these specifications. To enable realistic simulation of LLM training on heterogeneous infrastructure, we have identified the key abstractions "A" and components "C" that a heterogeneity-aware full-stack simulator must expose and implement. [A1] Device Group and Model Specification (§4.1). The simulator must support abstractions to ingest heterogeneous compute and AI model configurations. For each device group (DG), the specification must include information such as the PP stage, DP stage, number of model layers assigned, micro batch size, and the number of compute nodes (GPUs) in the device group. [A2] Heterogeneous Cluster and Host Topology Specification (§4.1). The simulator must ingest detailed network hardware specifications to accurately model performance variance. The abstraction should include the scale-up (host) and scale-out (cluster) topologies, along with link bandwidth and delays. The deployer can simulate various interconnect types, such as NVLink, PCIe, Ethernet, and InfiniBand, using appropriate topologies and bandwidth/delay configurations. [C1] Asymmetric Workload Generation (§4.2). Each device group profile may comprise a distinct number of model layers, tensor size, and micro-batch size. The simulator must generate a device-group-specific workload file (i.e., trace) and support Multiple-Instruction-Multiple-Data (MIMD) orchestration. [C2] Heterogeneity-aware Ring Construction (§4.3). The simulator must implement an algorithm that constructs an

4 𝑋𝑠𝑖𝑚 Simulator 𝑋𝑠𝑖𝑚 is a full-system, discrete-event training simulator for heterogeneous AI clusters. It extends the principles of the state-of-the-art simulator, SimAI [60], which is designed for training in homogeneous clusters. Figure 3 illustrates 𝑋𝑠𝑖𝑚’s key components distributed across three layers: (a) Input Specification: exposes abstractions to ingest heterogeneous cluster configurations generated by heterogeneous training optimization solutions, including multiple device groups (with varying compute, memory), topology (with scale-up and scale-out interconnects), model partitioning, and data partitioning, (b) Planning layer: prepares workload traces that reflect real-world characteristics and constructs heterogeneity-aware communication groups that would be used for gradient synchronization during simulation runtime, and (c) Simulation layer: uses an event-driven engine to model compute and collective operations (both scale-up and scale-out). Using the heterogeneity-aware framework parameters, i.e., per-device-group proto files, the planning layer constructs objects (DP/TP groups, PP/DP barrier groups, multi-ring perDP group) and mappings (model layer to chunk) necessary to simulate collective communication. 5

Figure 4. Heterogeneity-aware multi-ring LCM-based resharding in 𝑋𝑠𝑖𝑚 for non-uniform layer partitioning across high-compute (blue) and low-compute (red) device groups with layer-aware DP groups.

Figure 3. 𝑋𝑠𝑖𝑚’s Design The execution engine parses the workload trace and schedules compute events via a global priority-based event scheduler that supports concurrent per-device execution. For collective operations (e.g., AllReduce), it dynamically instantiates communication channels using the planning-phase configuration and invokes the selected network backend (NS-3 or htsim) through a send/recv interface. Our 𝑋𝑠𝑖𝑚 implementation spanned ∼ 8300 lines of code, including NS-3 backend implementation. 4.1

that can generate and map multiple workloads, rather than broadcasting a single static workload across the cluster. 𝑋𝑠𝑖𝑚’s Asymmetric Workload Generator (𝑆𝑈𝑇 𝑅𝐴𝐴𝑊 𝐺 ) ingests heterogeneity-aware framework and model parameters to produce distinct workload files for each device group, reflecting each DG’s unique computational and communication characteristics. 𝑆𝑈𝑇 𝑅𝐴𝐴𝑊 𝐺 (satisfies [C1]) operates in two phases: (a) In the capability profiling phase, 𝑆𝑈𝑇 𝑅𝐴𝐴𝑊 𝐺 , profiles per-device GPU compute capability (TFLOPS) for each device type in the heterogeneous cluster, using a sample GPU, and (b) In the trace generation phase, 𝑆𝑈𝑇 𝑅𝐴𝐴𝑊 𝐺 , extends SimAI’s framework interception mechanism. For each device group, the training framework is configured with DG-specific parameters such as assigned layers, micro-batch size, and tensor-parallel (TP) degree, and executed using the profiled compute capabilities. This execution generates a trace of computation and communication operations unique to the device group. These traces precisely capture the operator sequence, tensor shapes, and collective communication patterns executed during training.

Input Specification

To simulate training in heterogeneous clusters, the user provides as input (satisfies [A1] and [A2]): (1) heterogeneous cluster information, including compute, memory, cluster size, scale-up interconnect, and scale-out interconnect; (2) non-uniform workload to device mapping; and (3) model parameters. 𝑋𝑠𝑖𝑚 consumes deployment plans generated by heterogeneous training planners or specified manually by users. Its input abstraction captures the same logical parameters used by distributed training frameworks such as Megatron-LM [51] and DeepSpeed [47], including TP/DP/PP degrees, communication groups, pipeline stages, layer ranges, micro-batch sizes, and device placement. Heterogeneous deployments (e.g., HexiScale [65]) are represented through a protobuf-based specification that captures non-uniform configurations across device groups, such as GPU type, layer assignment, batch size, and TP degree. We adopt protobuf to align with production ML infrastructure practices (e.g., Ray [15]) and facilitate future integration with deployment pipelines. Network topology, bandwidth, and delay parameters are specified separately through a topology configuration file. Appendix §A provides example deployment and topology specifications. 4.2

4.3

Heterogeneity-aware Ring Construction.

Figure 4 shows an example that outlines the workflow for constructing objects and groups required for collective communication. The example shows mixed H100/A100 clusters using non-uniform layer partitioning. High-compute groups (blue) handle more layers (e.g., 20&15 layers, TP=3&2), while low-compute groups (red) handle fewer layers, balanced by non-uniform batches. In heterogeneous clusters, asymmetric pipeline partitioning can cause device groups to process overlapping but nonidentical layer ranges. A static DP group would therefore attempt to synchronize gradients for layers not present on all devices, leading to incorrect synchronization. In this section, we explain how 𝑋𝑠𝑖𝑚 addresses this challenge. Using the heterogeneity-aware framework parameters, i.e., per-device-group proto files, DP groups, and a Multi-ring per-DP group is constructed (satisfies [C2]). First, device groups are used to generate data parallel (DP) groups for

Asymmetric Workload Generator

Device groups may execute pipeline stages of vastly different depths (e.g., 20 vs. 12 layers), process different micro-batch sizes due to memory constraints, and use different tensorparallel (TP) degrees. This heterogeneity requires a simulator 6

Algorithm 1 Sweep-Line Algorithm for Dynamic Group Formation

Algorithm 2 LCM-Based Multi-Group Construction for Heterogeneous data Synchronization

1: Input: Set of Device Groups, D = {𝐷𝐺 1 , 𝐷𝐺 2 , . . . , 𝐷𝐺𝑛 } , where each 𝐷𝐺𝑖 has

1: Input: DP synchronization group, GDP , covering layer range [𝑙𝑠 , 𝑙𝑒 ] , with partici-

assigned layer range, [𝑠𝑖 , 𝑒𝑖 ] 2: Output: Set of DP synchronization groups, G𝐷𝑃 3: begin: 4: P ← ∅ ⊲ Stores boundary points for each DP 5: for each 𝐷𝐺𝑖 ∈ D do 6: P ← P ∪ {𝑠𝑖 , 𝑒𝑖 + 1} 7: end for 8: Psorted ← Sort ( P ) 9: Punique ← Deduplicate ( Psorted ) 10: G𝐷𝑃 ← ∅ 11: group_id ← 0 12: for 𝑖 = 0 to | Punique | − 2 do 13: seg start ← Punique [𝑖 ] 14: seg end ← Punique [𝑖 + 1] − 1 15: C←∅ ⊲ Holds set of DGs that cover [𝑠𝑒𝑔𝑠𝑡𝑎𝑟𝑡 , 𝑠𝑒𝑔𝑒𝑛𝑑 ] 16: for each 𝐷𝐺 𝑗 ∈ D do 17: if 𝑠 𝑗 ≤ seg start and 𝑒 𝑗 ≥ seg end then 18: C ← C ∪ {𝐷𝐺 𝑗 } 19: end if 20: end for 21: if | C | ≥ 2Ðthen 22: R ← 𝐷𝐺 𝑗 ∈C GetRanks (𝐷𝐺 𝑗 ) 23: 𝐺 ← CreateDPGroup (group_id, [seg start , seg end ], R ) 24: G𝐷𝑃 ← G𝐷𝑃 ∪ {𝐺 } 25: group_id ← group_id + 1 26: end if 27: end for 28: return G𝐷𝑃

pating device groups, {𝐷𝐺 1 , 𝐷𝐺 2 , . . . , 𝐷𝐺𝑘 } , having TP degrees, {𝑡 1 , 𝑡 2 , . . . , 𝑡𝑘 }

2: Output: R , Multi-ring communication topology for gradient synchronization 3: begin: 4: 𝐿 ← lcm (𝑡 1 , 𝑡 2 , . . . , 𝑡𝑘 ) 5: R ← ∅ ⊲ Set of communication rings in a DP group 6: for 𝑐 = 0 to 𝐿 − 1 do ⊲ Create one ring per LCM chunk 7: participants𝑐 ← ∅ 8: for each 𝐷𝐺𝑖 ∈ {𝐷𝐺 1 , 𝐷𝐺 2 , . . . , 𝐷𝐺𝑘 } do 9: 𝑡𝑖 ← TP_degree (𝐷𝐺𝑖 ) 10: ranks𝑖 ← GetRanks (𝐷𝐺𝑖 ) 11: for each rank 𝑟 ∈ ranks𝑖 do ⊲ Mapping chunk to a set of local ranks 12: local_rank ← 𝑟 mod 𝑡𝑖 13: if (𝑐 mod 𝑡𝑖 = local_rank) then 14: participants𝑐 ← participants𝑐 ∪ {𝑟 } 15: end if 16: end for 17: end for 18: Ring𝑐 ← CreateCommRing (𝑐, participants𝑐 ) 19: R ← R ∪ { Ring𝑐 } 20: end for 21: return R

The sweep-line algorithm runs once during the planning phase, and it requires 𝑂 (𝐷 log 𝐷 + 𝐷 · 𝑆) time, where 𝐷 is the number of DGs and 𝑆 is the number of unique layer segments (typically 𝑆 ≤ 2𝐷). In realistic deployments, pipeline partitions do not exceed ≈ 8 stages [36, 37], and many DGs share identical layer boundaries, which keeps 𝑆 small. As a result, the sweep-line construction scales efficiently in practice, with 𝑆 bounded by a practical and small constant making the 𝑂 (𝐷 · 𝑆) term negligible even at cluster scale. LCM-based Multi-Ring construction algorithm. The heterogeneity-aware ring topology is constructed automatically during the planning phase. For each DP synchronization group identified by the Sweep-Line algorithm, 𝑋𝑠𝑖𝑚’s LCM-based Multi-Ring Construction algorithm (see Algorithm 2) performs the following steps: (1) Generate chunk-to-rank mappings by determining which ranks (across all participating DGs) are responsible for each chunk index, 𝑐 ∈ [0, 𝐿 − 1], using the interleaved assignment formula, (2) instantiate communication rings by creating a MockNccl [60] communicator for each chunk 𝑐 containing only the ranks assigned to that chunk, configured with the appropriate network topology (e.g., ring, tree, or hierarchical) based on the underlying interconnect structure, (3) The execution engine stores these mappings in a layer-aware routing table indexed by (𝑙𝑎𝑦𝑒𝑟 _𝑛𝑢𝑚𝑏𝑒𝑟, 𝑐ℎ𝑢𝑛𝑘_𝑖𝑛𝑑𝑒𝑥), enabling the simulator to quickly look up the correct communication ring during the backward pass. The LCM-based synchronization mechanism is a critical innovation that enables 𝑋𝑠𝑖𝑚 to accurately model the gradient aggregation patterns of state-of-the-art heterogeneous training frameworks such as HexiScale and Metis, which rely on asymmetric TP configurations to balance compute and memory constraints across mixed-generation GPU clusters.

gradient synchronization using our Sweep-line algorithm (see Algorithm 1). We have designed an LCM-based Multiring construction and an LCM-based Gradient chunking algorithm that generates the multi-ring topology and chunk sizing/mapping for each DP group. In Appendix §B, we illustrate the algorithms using a concrete example. Sweep-line algorithm [28]. Consider four Device Groups with the following assigned layer ranges [𝑠𝑖 , 𝑒𝑖 ]: 𝐷𝐺 0 with layer range [1,20] and ranks {0, 1, 2}, 𝐷𝐺 1 with layer range [21,32] and ranks {5, 6}, 𝐷𝐺 2 with layer range [1,15] and ranks {3, 4}, and 𝐷𝐺 3 with layer range [16,32] and ranks {7, 8, 9}. For each Device Group, the sweep-line algorithm first collects all layer boundaries, i.e., {1, 21, 21, 33, 1, 16, 16, 33} (for algorithmic optimizations, 𝑒𝑖 is incremented by one). These boundaries are sorted, and duplicates are removed, resulting in the set {1, 16, 21, 33}, which defines three unique segments: [1,15], [16,20], and [21,32]. For each segment, the algorithm identifies the Device Groups that fully cover the segment and forms a DP group, GD P . Segment [1,15] is covered by both 𝐷𝐺 0 and 𝐷𝐺 2 , so a DP synchronization group is formed that contains the union of their ranks which results in {0, 1, 2, 3, 4}. Similarly, segment [16,20] is covered by 𝐷𝐺 0 and 𝐷𝐺 3 , forming a second DP group with ranks {0, 1, 2, 7, 8, 9} and lastly, segment [21,32] is covered by 𝐷𝐺 1 and 𝐷𝐺 3 , which results in a third DP group with ranks {5, 6, 7, 8, 9}. This decomposition enables Layer-Aware Multi-Group Membership: rank 0 (from 𝐷𝐺 0 ) participates in two different DP groups, one for synchronizing layers 1–15 and another for layers 16–20, depending on which layer’s gradients are being processed. 7

Algorithm 3 LCM-based Gradient/Data Chunking 1: Input: DP synchronization group, GDP , with participating device groups,

{𝐷𝐺 1 , 𝐷𝐺 2 , . . . , 𝐷𝐺𝑘 } , having TP degrees, {𝑡 1 , 𝑡 2 , . . . , 𝑡𝑘 } , with communication volume, 𝑑 2: Output: For each 𝐷𝐺𝑖 ∈ {𝐷𝐺 1 , 𝐷𝐺 2 , . . . , 𝐷𝐺𝑘 } per-chunk per-rank volume 3: 𝐿 ← lcm(𝑡 1 , 𝑡 2 , . . . , 𝑡𝑘 ) 4: for each 𝐷𝐺𝑖 ∈ {𝐷𝐺 1 , 𝐷𝐺 2 , . . . , 𝐷𝐺𝑘 } do 𝑑 5: data_per_rank𝐷𝐺𝑖 ← 𝑡𝑖 𝐿 6: chunk_multiplier𝐷𝐺𝑖 ← 𝑡𝑖 data_per_rank𝐷𝐺𝑖 7: data_per_chunk_per_rank ← chunk_multiplier𝐷𝐺𝑖 8: end for

4.4

Figure 5. Overview of contributions to htsim implementation. communication groups and chunking derived using the LCMbased synchronization algorithms. After local computation completes, devices in the DP Barrier Group (DBG) trigger gradient synchronization across data-parallel replicas. Together, the PBG and DBG abstractions enable faithful simulation of pipeline execution, interstage communication, and heterogeneity-induced idle time. 𝑋𝑠𝑖𝑚 provides reusable abstractions for emerging pipelineparallel algorithms. Beyond the current GPipe implementation, we are extending support to 1F1B [24] and other pipeline schedules.

Heterogeneity-aware Chunk Partitioning

With non-uniform workload partitioning, the gradient tensor shapes become incompatible for aggregation. 𝑋𝑠𝑖𝑚 designs an LCM-based Gradient Chunking algorithm (see Algorithm 3) that repartitions the gradients to establish a common synchronization granularity across the communicating heterogeneous DGs in a DP group. (satisfies [C3]) In Appendix §C, we illustrate the algorithms using a concrete example. The core insight is that if two DGs participating in the same DP group have TP degrees 𝑡𝑖 and 𝑡 𝑗 , their gradient chunks can be made compatible by subdividing them into a finer granularity equal to the Least Common Multiple (LCM) of the two TP sizes. For a DP group spanning layers [𝑙𝑠 , 𝑙𝑒 ], with TP degrees 𝑡 1, 𝑡 2, . . . , 𝑡𝑘 , we set L = lcm(𝑡 1, 𝑡 2, . . . , 𝑡𝑘 ). Given a gradient communication volume, 𝑑, for the DP group, each device group, 𝐷𝐺𝑖 , computes its per-rank gradient volume as 𝑡𝑑𝑖 , and each rank in 𝐷𝐺𝑖 further subdivides the volume into equal-sized chunks, 𝑡𝐿𝑖 . Hence, each rank contributes a gradient chunk of size 𝑑𝐿 units to each ring it participates in. This guarantees that all communication rings operate on identically sized gradient chunks, ensuring correctness. Algorithm 3 records these per-chunk, per-rank communication volumes, which are then assigned to ranks using an interleaved strategy, and dedicated communication rings are constructed for each chunk. We discuss the bounds and tradeoffs of 𝑋𝑠𝑖𝑚’s LCM-based chunking technique in §E. 4.5

4.6

Network Simulation

Accurate performance prediction for heterogeneous LLM training requires modeling the interaction between collective communication and diverse network fabrics. Packetlevel simulators such as NS-3 offer high fidelity by modeling per-packet behavior and protocol state, but their cost is prohibitive at scale. A single AllReduce can generate millions of packets, making packet-level simulation of even one training iteration take hours. For example, in our evaluation, simulating one Llama 7B training iteration on 128 GPUs took 2 hours (Figure 16). This limits exploration of heterogeneous deployment and network design choices. 𝑋𝑠𝑖𝑚 addresses this scalability gap with dual network backends: NS-3 for protocol-accurate analysis and htsim for fast, scalable simulation. This design lets users trade fidelity for speed, depending on their experimental goals (satisfies [C5]). Extensions to NS-3 for Heterogeneous Clusters. NS3 models the cluster network topology [4]. 𝑋𝑠𝑖𝑚 extends NS-3 to capture host-network topology, including scale-up interconnects, PCIe switches, and NIC processing. Packets traverse a detailed path comprising PCIe switches, NICs, Topof-Rack (ToR) switches, and scale-up and core fabric links, each instantiated with realistic bandwidths and latencies. This model captures mixed-generation deployments (e.g., ConnectX-6 vs. ConnectX-7, PCIe Gen4 vs. Gen5, varying ToR uplink rates), enabling 𝑋𝑠𝑖𝑚 to quantify how slower NICs, additional PCIe hops, or reduced bandwidth introduce stragglers and inflate collective completion times. Our NS-3 implementation spanned ∼ 380 lines of code. Extensions to htsim for Heterogeneous Clusters. 𝑋𝑠𝑖𝑚 integrates htsim [14] as a network backend that models communication at a flow-level abstraction for speedup. We design and implement three extensions to htsim to enable simulation of heterogeneous host network topology.

Pipeline-Parallel Barrier and Communication Modeling

𝑋𝑠𝑖𝑚 models stragglers arising from non-uniform workload distribution across pipeline stages satisfies ([C4]) and supports hybrid PP–TP–DP execution. To enforce pipeline dependencies, 𝑋𝑠𝑖𝑚 creates a Pipeline Barrier Group (PBG)5 , ensuring that stages within a model replica execute in order and that later stages wait for preceding stages to complete. Note that the SOTA simulators, such as SimAI [60], ignore this. 𝑋𝑠𝑖𝑚 further inserts explicit point-to-point activation and gradient transfers between adjacent pipeline stages, with 5 Pipeline Barrier Group (PBG): The set of ranks within a single DP stage.

8

(1) PCIe switch model. In the htsim [14] design, GPUs are directly connected to the ToR switches. To simulate PCIe switching cost, we implemented an additional layer of switches (using htsim’s FatTreeSwitch class) between the GPUs and ToR, which acts as the NIC (see Figure 5). GPU– NIC and NIC–ToR are implemented as one-to-one and oneto-many connections, respectively. GPU–NIC interconnect is configured with real-world bandwidth and processing delays to simulate PCIe costs for heterogeneous nodes, while the remainder of the scale-out topology remains unchanged. (2) Scale-up network model. To model intra-node, interGPU communication, we designed and implemented a dedicated scale-up network layer (see Figure 5) (using FatTreeSwitch htsim class). To represent the scale-up fabric, the GPUs connect to a scale-up switch configured with low latency (sub-µs) and high bandwidth (say, 600–900 GB/s). We implemented the htsim function, get_bidir_paths() that updates routing tables such that GPUs within the same machine connect to the same scale-up switch, bypassing ToR and Aggregator (AGG) switches for TP and intra-node PP communication. The scale-up switch models scale-up non-blocking behavior with negligible queueing delay. (3) Rail-Optimized topology integration for training clusters. The current htsim implementation does not support the rail-optimized topology commonly used in AI training clusters. htsim models a fat-tree topology, which restricts the network to ToR and aggregation switches with full all-toall connectivity. To isolate bandwidth across collective traffic, a rail-optimized topology connects GPUs with identical local ranks to the same ToR via dedicated NICs, forming isolated rails that bypass aggregation switches during collective operations such as AllReduce. We extend htsim to model rail-optimized topology by partitioning the network graph into isolated rail subgraphs. To model this scale-out network, we implemented the htsim function, get_bidir_paths(), which updates the routing table for both the scale-out and scale-up networks. This design scales to large dataplane elements while accurately capturing rail isolation, where GPUs of the same local rank share dedicated ToR channels. Our htsim implementation spanned ∼ 500 lines of code. Our htsim prototype can independently be integrated with trace-driven simulators such as ATLAHS [50] that support generic workloads, i.e., AI, as well as HPC

5

Table 4. Cluster Configuration. Column 2 shows the cluster size and GPU layout (e.g., 2 × (4 × H100) means two nodes with four H100 GPUs each). Column 3 shows model parallelism strategy (e.g., 4 × ((4 × TP)-DP) means four dataparallel replicas, each using 4-way tensor parallelism) Config C1 C2 C3 C4 C5 C6 C7 C8 C9 C10 C11 C12 C13 C14 C15 C16

Compute Cluster 2 × H100 2 × A100 2 × (4 × H100) 2 × (4 × A100) 2 × (4 × H100) 2 × (4 × A100) 4 × (4 × H100) 4 × (4 × A100) 1 × A100, 1 × H100 2 × A100, 2 × H100 2 × A100, 2 × H100 2 × A100, 2 × H100 4 × A100, 4 × H100 4 × A100, 4 × H100 4 × A100, 4 × H100 2 × (4 × A100), 2 × (4 × H100)

Model Distribution DP DP DP DP 2 × ((4 × TP) - DP) 2 × ((4 × TP) - DP) 4 × ((4 × TP) - DP) 4 × ((4 × TP) - DP) DP A (2 × DP) - DP - H (2 × DP) A (2 × TP) - DP - H (2 × TP) A (2 : 2 PP) - DP - H (2 : 2 PP) A (4 × DP) - DP - H (4 × DP) A (4 × TP) - DP - H (4 × TP) A (3 × TP : 1) PP - DP - H (3 × TP : 1) PP 2 × (A (4 × TP) - DP - H (4 × TP))

Table 5. Real-world Cluster Interconnect Configuration. GPU A100 40GB [41] H100 80GB [42] H200 141GB [13] B200 192GB [12]

Scale-Up BW 2400 Gb/s 3600 Gb/s 3600 Gb/s 7200 Gb/s

Scale-Up Delay 30.66 ns 20.44 ns 20.44 ns 10.22 ns

PCIe BW 256 Gb/s 512 Gb/s 512 Gb/s 512 Gb/s

PCIe Delay 2×287.5 ns 2x143.75 ns 2x143.75 ns 2x143.75 ns

NIC BW

NIC delay

10 Gb/s 10 Gb/s 200 Gb/s 200 Gb/s

3 µs 1.5µs 368 ns 368 ns

and H100 GPUs to run simulations for homogeneous and heterogeneous configurations. All experiments were run for one training iteration, unless mentioned otherwise. Table 6 shows the heterogeneous cluster’s compute and interconnect specifications, which we used for simulation-only experiments. Real-world testbed. We evaluate 𝑋𝑠𝑖𝑚 on a heterogeneous testbed comprising: (i) two AMD EPYC 9354 nodes, each with 2×A100 and 4×H100 GPUs, scale-up communication uses PCIe Gen4/Gen5 interconnects, while scale-out communication uses Intel X710 and X550T NICs. (ii) an Intel Xeon 8480C node with 8×H200 GPUs, and an Intel Xeon 8570 node with 8×B200 GPUs, with NVLink for scale-up and ConnectX-7 for scale-out, as summarised in Table 5. The nodes are distributed across racks connected through ToR and aggregation switches. Results are averaged over three runs of a single training iteration unless stated otherwise. Unless specified, training-time prediction experiments use the NS-3 backend. Metrics. We measure (1) Prediction accuracy as the percentage error between the simulated training time (𝑋𝑠𝑖𝑚, SimAI [60], or Hexiscale [65]) and real hardware execution time, percentage error of scale-up (TP communication) and scale-out (DP communication using multi-ring) implementation with a real cluster (2) Simulation runtime as the wall-clock time required to complete a full training iteration simulation, (3) GPU idle time (a.k.a., Straggler waiting time) measures the GPU waiting time due to AllReduce

Evaluation

We run Llama 2 (7B, 13B) and GPT-175B models across experiments, with fixed optimizers and controlled DP/TP/PP settings to isolate each evaluation dimension. Simulation testbed. We run 𝑋𝑠𝑖𝑚 and SimAI simulation frameworks on two 32-core AMD EPYC-9354 nodes with 2 × 𝐴100 and 4 × 𝐻 100 GPUs. We generate model workloads using the AICB workload benchmark [20] on A100 9

Table 6. Cluster configurations used for simulation-only experiments. Ampere [10, 41] Hopper [10, 42]

A100 (40GB) H100 (80GB)

Llama 7B

79.80%

4.66%

Training time per iteration (sec)

30 79.15% 25 20 15 10 5 0 SimAI C2

NVLink BW (Gbps) 2400 (Gen 3) 3600 (Gen 4)

NVLink delay (ns) 30.66 20.44 Llama 7B

PCIe BW (Gbps) 256 (Gen 4) 512 (Gen 5)

Hexiscale

PCIe latency (ns) 2 × 287.5 2 × 143.75

Real

104 103 Llama 13B

Figure 6. Training time per iteration for Llama 7B on a heterogeneous cluster shows that 𝑋𝑠𝑖𝑚 closely matches real hardware with <5% error, while SimAI incurs large errors due to explicit heterogeneity modeling.

103 C9

C11 C12 Cluster Configuration

C10

Figure 7. Training time per iteration for Llama 7B and 13B across heterogeneous cluster configurations, comparing Hexiscale, 𝑋𝑠𝑖𝑚, and real hardware.

completion skew across DP groups in a heterogeneous cluster, and (4) Total Cost of Ownership (TCO) represents capital expenditure on GPU hardware, and it is computed as TCO = CapEx/𝑇 𝑟𝑎𝑖𝑛𝑖𝑛𝑔𝑇 𝑖𝑚𝑒, where CapEx is the GPU cost in USD, with the unit, “$/GPU-hour", (5) Pipeline bubble time is the GPU waiting time during pipeline resharding and communication. Evaluation Questions. (1) Can we use existing training simulators that assume homogeneous AI clusters to simulate heterogeneous workload? (2) How accurately does 𝑋𝑠𝑖𝑚 predict real-world training performance on heterogeneous clusters compared to existing heterogeneity-aware simulators? (3) How does the scalability of htsim versus NS-3 affect simulation time in 𝑋𝑠𝑖𝑚? (4) How accurately does 𝑋𝑠𝑖𝑚 predict collective communication overheads across scale-up fabric, i.e., NVLink? (5) How accurately does 𝑋𝑠𝑖𝑚 predict the time taken for Data Parallelism (DP) multi-ring collective communication operations across different model scales on a heterogeneous distributed cluster? (6) What is the source of prediction error in 𝑋𝑠𝑖𝑚? (7) How does 𝑋𝑠𝑖𝑚 perform compared to the SOTA resharding solutions, Hetauto and AlpaComm, with varying topology abstractions? (8) How does 𝑋𝑠𝑖𝑚 compare with state-of-the-art training simulators for homogeneous AI clusters? (with respect to prediction accuracy and simulation runtime) (§D)

NIC processing delay (ns) 368 368

Xsim

104

SimAI Xsim Real C1 C9 C9 Cluster Configuration

NIC BW (Gbps) 400 (ConnectX-7) 400 (ConnectX-7)

Network Simulation runtime (per AllReduce) (ms)

GPU

Training time per iteration (ms) (in log scale)

Architecture

htsim

108 AllReduce, 1Mb size 106 104 102

13X

29X

32X

35X

NS-3 38X

41X

108 AllReduce, 4Mb size 106 104 102

16X

8

31X

16

31X

32

30X

64

34X

128

Cluster Size

38X

256

44X

44X

512

44X

47X

1024

Figure 8. Cluster size vs. Network simulation runtime. NS-3’s protocol execution incurs higher overhead at larger scales and smaller message sizes (16×–47×).

(9) How well does 𝑋𝑠𝑖𝑚 scale with model and cluster size? (§D) (10) How effectively does 𝑋𝑠𝑖𝑚 support system and capacity planning by exposing heterogeneity-aware metrics such as GPU idle time (due to stragglers) and Total Cost of Ownership (TCO) across cluster designs?(§D) (1) Comparing heterogeneous deployment’s training time in the wild (i.e., the ground truth), with SimAI and 𝑋𝑠𝑖𝑚. We run the heterogeneous deployment, 𝐶9 (see Table 4 and Fig. 6), i.e., 1 × A100, 1 × H100, over the realworld cluster, and also simulate it using 𝑋𝑠𝑖𝑚. Since SimAI does not support heterogeneous system architecture, we deploy 𝐶9 over two homogeneous configurations, 𝐶1 (2 × H100) and 𝐶2 (2 × A100). Results show that SimAI incurs large prediction errors ( 80%). This is because SimAI assumes homogeneous AI clusters and cannot accurately simulate collective communication for non-uniform model parallelism, device bandwidth, delays, and communication asymmetry. 𝑋𝑠𝑖𝑚, by explicitly modeling heterogeneity and workload transformation, achieves near-real accuracy (<5% error). (2) Comparing heterogeneous deployment’s training time in the wild (i.e., the ground truth), with Hexiscale ’ s [65] heterogeneity-aware analytical simulator, and 𝑋𝑠𝑖𝑚. Fig. 7 demonstrates the training time for Hexiscale, 𝑋𝑠𝑖𝑚, and real-world evaluations, across four heterogeneous configurations, 𝐶9, 𝐶10, 𝐶11, and 𝐶12 (see Table 4). These configurations comprise heterogeneous compute, interconnect, and non-uniform hybrid parallelism strategies. 10

Llama 450B

layernorm attn layernorm attn layernorm attn layernorm attn layernorm attn & mlp & mlp & mlp & mlp & mlp

Layers for scale-up communication

6.2%

8.4%

4.3% 6.2%

9.7% 11.1% 10.5% 10.9% 17.1%

grad_ grad_ grad_ grad_ grad_ grad_ grad_ grad_ grad_ grad_ gather param gather param gather param gather param gather param

Layers for DP communication

Llama 450B

14.2%

GPT 175B

Llama 70B

109

108

105

Llama 13B

Xsim Real

Llama 7B

Collective Communication Time (ns)

GPT 175B

Llama 70B

Llama 13B

Llama 7B

Collective Communication Time (ns)

106

Real Xsim ns3 Xsim htsim

grad_gather grad_param_comm optimizer

layernorm embedding_grads embedding_layer

attention_layer mlp_layer cross_entropy

B200 x 4 & H200 x 8

GroundTruth Xsim

B200 x 2 & H200 x 8

GroundTruth Xsim

B200 x 2 & H200 x 4

109

1011 1010 Collective Communication Time (ns)

1012

GroundTruth Xsim

Figure 9. Comparison of isolated scale-up Figure 10. Comparison of scale-out DP col- Figure 11. Layer-wise validation of execution TP collective communication time between lective communication time between 𝑋𝑠𝑖𝑚 time across diverse heterogeneous cluster con𝑋𝑠𝑖𝑚 (using NS-3 and htsim backends) and and a real cluster making the multi-ring on figurations. The breakdown isolates of layers, real cluster of 8× H200 NVLink nodes with a non-isolated cluster of 4× H100 & 2× A100 demonstrating 𝑋𝑠𝑖𝑚’s precision down to individan average error of 5.5%. cluster. ual layers

For both Llama 7B and Llama 13B, we observe that Hexiscale underestimates training time by 90–99%, since its analytical model does not account for the synchronization overheads introduced by asymmetric communication during collective operations. In contrast, 𝑋𝑠𝑖𝑚 consistently achieves near-real accuracy, with <5% error in most configurations, demonstrating its effectiveness as a full-stack heterogeneity-aware training simulator. This experiment highlights the fact that network communication cannot be ignored in large-scale AI training clusters. (3) Network simulation scalability: NS-3 vs. htsim. Fig. 8 shows how 𝑋𝑠𝑖𝑚’s network simulator backends, NS-3 [4] (uses 8 CPU cores) and htsim (uses 1 CPU), scale with increasing cluster size. We evaluate the simulators over the basic communication unit in the AI training cluster, i.e., AllReduce, with two sizes, 1𝑀𝑏 and 4𝑀𝑏. NS-3 provides packet-level fidelity with protocol execution, hence its simulation overhead increases with cluster size, particularly under small chunk sizes and large data volumes. htsim is 16×–44× faster than NS-3 in the cluster from 8 nodes to 1024 nodes, since htsim uses flow-level abstraction and abstracts away the packet header parsing. A simulator user interested in protocol-level analysis must use NS-3. However, network traffic analysis, such as congestion control, and microarchitectural analysis, such as scaleup switch queueing delays, can be analyzed using htsim. (4) How accurately does 𝑋𝑠𝑖𝑚 predict scale-up collective communication overhead? To evaluate the predictive fidelity of 𝑋𝑠𝑖𝑚’s scale-up network models, we measure TP collective communication time on an 8×H200 NVLink node across models ranging from Llama-7B to Llama-450B and GPT-175B. Figure 9 compares the NS-3 and htsim backends against real hardware measurements for LayerNorm and Attention/MLP collectives. For the large-message Attention/MLP collectives, both backends closely track hardware performance, with average error below 5.5%. The small remaining error is primarily attributable to NCCL runtime overheads and htsim’s flow-level abstraction. These results validate 𝑋𝑠𝑖𝑚’s modeling of scale-up communication while

demonstrating that htsim provides a scalable alternative when packet-level fidelity is not required. (5) How accurately does 𝑋𝑠𝑖𝑚 predict the time taken for Data Parallelism (DP) multi-ring collective communication operations across different model scales on a heterogeneous distributed cluster? To evaluate 𝑋𝑠𝑖𝑚’s multi-ring chunk-partitioning algorithm in a scale-out setting, we measure DP synchronization time on a heterogeneous cluster comprising 4×H100 and 2×A100 GPUs across models ranging from Llama-7B to Llama-450B and GPT-175B. Figure 10 compares simulated and real execution times for gradient-gather and gradient-parameter collectives. 𝑋𝑠𝑖𝑚 closely tracks hardware performance across all models, with error decreasing as model size increases. The remaining discrepancy at large scales stems from NCCL implementation constraints: while 𝑋𝑠𝑖𝑚 abstracts multi-ring communication as fully parallel chunk transfers, real NCCL executions are limited by link assignment and resource contention, which can serialize communication across rings [11]. (6) What is 𝑋𝑠𝑖𝑚’s source of error in the prediction of execution and communication time compared to realworld execution times? – a layer-wise Ablation study. To evaluate 𝑋𝑠𝑖𝑚’s fine-grained fidelity, we compare layerwise execution times across three heterogeneous clusters: 𝐵200 × 4 & 𝐻 200 × 8, 𝐵200 × 2 & 𝐻 200 × 4, and 𝐵200 × 2 & 𝐻 200 × 8. Figure 11 shows that system closely matches measured execution times across compute, communication, and optimizer operations spanning three orders of magnitude. Compute-intensive layers such as attention_layer and mlp_layer exhibit less than 1% error, while communication operations (grad_gather and grad_param_comm) remain within 1–4% error across all hardware configurations. These results demonstrate that 𝑋𝑠𝑖𝑚 accurately captures both computation and heterogeneous scale-out communication behavior at layer granularity. (7) How does 𝑋𝑠𝑖𝑚 perform compared to the SOTA resharding solutions with varying topology abstractions? We compare 𝑋𝑠𝑖𝑚’s predicted performance with two stateof-the-art resharding schemes, AlpaComm [75] and HetAuto [44], across heterogeneous deployments (recall §2.4): 11

Hetauto

Xsim

Alpa-Comm

HeteCCL [25] and ForestColl [73] design efficient scheduling solutions. 𝑋𝑠𝑖𝑚 acts as a bridge between heterogeneous compute deployment works, such as Metis [58], and these works. Resharding Across Heterogeneous Tensor Layouts. HetAuto [44] addresses tensor-layout mismatches using a GCD-based gather–P2P–scatter resharding strategy across heterogeneous pipeline stages. AlpaComm [75] uses slice-based decomposition with non-uniform, point-to-point sender–receiver mappings, while NTP [2] reconstructs temporary uniform layouts through point-to-point transfers before invoking standard collectives. In contrast, 𝑋𝑠𝑖𝑚 employs LCM-based chunk partitioning and multi-ring collectives, enabling balanced synchronization across heterogeneous TP/DP/PP configurations without relying on sequential resharding phases. Full-stack simulators for distributed training. Full-stack AI training simulators such as ASTRA-sim [62], SimAI [60], Multiverse [21], Phantora [46], and ATLAHS [50] model both system and network behavior. ATLAHS targets generalpurpose distributed applications rather than AI training, while ASTRA-sim and Meta’s Arcadia [16] rely on production traces and offer limited support for heterogeneous communication patterns. Phantora [46] executes unmodified ML frameworks in a distributed containerized environment but does not model compute or network heterogeneity. Analytical full-stack simulators. Analytical simulators such as Echo [17], Neusight [29], and vTrain [5] prioritize simulation speed over fidelity. However, these approaches trade modeling depth for speed and do not capture complex heterogeneous collective behavior.

Training Total Pipeline Bubble Total Time (ms) Time (ms)

600 400 200 0 60 40 20 0

H100 x 6 - A100 x 4

H100 x 8 - A100 x 1 H100 x 4 - A100 x 4 Experimental Configurations (Topology Pairs)

Figure 12. Total training time & exposed PP communication time comparison across asymmetric topology pairs comparing different re-sharding technique algorithms. 𝐻 100 × 6 → 𝐴100 × 4, 𝐻 100 × 8 → 𝐴100 × 1, and 𝐻 100 × 4 → 𝐴100 × 4. Figure 12 reports total training time and exposed pipeline bubble time. HetAuto incurs higher training time in asymmetric configurations due to its hierarchical gather–P2P–scatter workflow, whose benefits diminish as the GCD between source and destination parallelism degrees decreases. AlpaComm and 𝑋𝑠𝑖𝑚 avoid these additional aggregation phases through point-to-point communication, reducing end-to-end training time. For symmetric configurations, all three schemes achieve similar performance. For pipeline bubble time, 𝑋𝑠𝑖𝑚 and HetAuto partition tensors into balanced communication units, resulting in lower stage imbalance and idle time. In contrast, AlpaComm’s cutpoint-union approach produces non-uniform communication slices, leading to load imbalance, communication skew, and higher pipeline bubble time.

6

Related work

Optimizations in the presence of heterogeneous distributed training infrastructure. Prior heterogeneous AI training systems [27, 43, 58, 64, 65, 67, 70, 71] optimize hybrid parallelism and workload placement across heterogeneous compute and network resources. Learning-based [31, 68] and analytical planners [33, 55] further automate deployment decisions, but abstract away detailed communication and network dynamics. 𝑋𝑠𝑖𝑚 complements these works by enabling full-stack simulation of heterogeneous compute, communication, and network configurations. Heterogeneity-aware simulators. Heterogeneity-aware simulators such as LLMServingSim [6] and HeteroSim [69] focus on inference or assume uniform workload partitioning, limiting support for heterogeneous training. Similarly, Wu et al. [63] use SimAI [60] for heterogeneous planning without extending it to model heterogeneous workload assignment or collective communication. In contrast, 𝑋𝑠𝑖𝑚 supports non-uniform workload partitioning, heterogeneous synchronization, and full-stack training simulation. Heterogeneity-aware collective communication To address cluster topology and link bandwidth heterogeneity,

7

Conclusion

𝑋𝑠𝑖𝑚 enables faithful simulation of heterogeneous hybridparallel training through device groups, sweep-line DP formation, LCM-based synchronization, and gradient resharding. 𝑋𝑠𝑖𝑚 ’s extensible abstractions support emerging tensorresharding and pipeline-parallel techniques, while dual NS-3 and htsim backends provide a flexible trade-off between network fidelity and simulation scalability. We position 𝑋𝑠𝑖𝑚 as a foundational step towards a unified evaluation platform for heterogeneous AI training. As models, hardware, and parallelization strategies co-evolve, 𝑋𝑠𝑖𝑚 equips the community to systematically study their interactions and reason about future training infrastructures before deployment at scale.

References [1] Meta AI. 2025. The LLaMA 4 Herd: The Beginning of a New Era of Natively Multimodal AI Innovation. https://ai.meta.com/blog/ llama-4-multimodal-intelligence/ Accessed from Meta AI Blog on 2026-02-06. 12

[2] Daiyaan Arfeen, Dheevatsa Mudigere, Ankit More, Bhargava Gopireddy, Ahmet Inci, and Gregory R Ganger. 2025. Nonuniformtensor-parallelism: Mitigating gpu failure impact for scaled-up llm training. arXiv preprint arXiv:2504.06095 (2025). [3] ascentoptics. 2025. InfiniBand vs RoCE network fabrics: RDMA interconnect comparisons. https://ascentoptics.com/blog/infiniband-vsroce-which-is-better-suited-for-ai-data-center-networks/ InfiniBand delivers engineered lossless RDMA fabrics with ultra-low latency and high throughput (up to 400 800 Gbps+ per port), while RoCE (especially RoCEv2) brings RDMA to Ethernet with high bandwidth and low latency under loss-free configurations. Accessed: 2026-02-06. [4] Songyuan Bai, Hao Zheng, Chen Tian, Xiaoliang Wang, Chang Liu, Xin Jin, Fu Xiao, Qiao Xiang, Wanchun Dou, and Guihai Chen. 2024. Unison: a parallel-efficient and user-transparent network simulation kernel. In Proceedings of the Nineteenth European Conference on Computer Systems. 115–131. [5] 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 2024 57th IEEE/ACM International Symposium on Microarchitecture (MICRO). IEEE, 153–167. [6] Jaehong Cho, Minsu Kim, Hyunmin Choi, Guseul Heo, and Jongse Park. 2024. LLMServingSim: A HW/SW Co-Simulation Infrastructure for LLM Inference Serving at Scale. In 2024 IEEE International Symposium on Workload Characterization (IISWC). IEEE, 15–29. [7] Alibaba Cloud. 2021. Apsara Conference 2021 | Alibaba Cloud Released the Fourth-Generation SHENLONG Architecture. https:// www.alibabacloud.com/blog/598193 Accessed: 2026-02-06. [8] Alibaba Cloud. 2026. Heterogeneous Computing. https://www.alibabacloud.com/en/product/ heterogeneous_computing?_p_lc Accessed: 2026-02-06. [9] Google Cloud. 2026. GPU networking overview. https: //docs.cloud.google.com/ai-hypercomputer/docs/networkingoverview Accessed: 2026-02-06. [10] NVIDIA Corporation. ConnectX-7 Datasheet. https: //www.nvidia.com/content/dam/en-zz/Solutions/networking/ ethernet-adapters/connectx-7-datasheet-Final.pdf. Accessed: 2026-02-06. NCCL Developer Guide: Collec[11] NVIDIA Corporation. 2018. tive Communication Primitives (Version 2.0.5). NVIDIA Corporation. https://docs.nvidia.com/deeplearning/nccl/archives/nccl_205/ nccl-developer-guide/index.html [12] NVIDIA Corporation. 2026. NVIDIA DGX B200: The Foundation for Your AI Factory. https://www.nvidia.com/en-in/data-center/dgxb200/. Accessed: 2026-06-11. [13] NVIDIA Corporation. 2026. NVIDIA H200 GPU. https:// www.nvidia.com/en-in/data-center/h200/. Accessed: 2026-06-10. [14] Broadcom CSG. HTSim Network Simulator. https://github.com/ Broadcom/csg-htsim Accessed from GitHub repository. Accessed: 2026-02-06. [15] Ray Documentation. 2025. Set Up a gRPC Service — Ray Serve gRPC Guide. https://docs.ray.io/en/latest/serve/advanced-guides/grpcguide.html Accessed: 7 Feb 2026. [16] Meta Engineering. 2023. Arcadia: An end-to-end AI system performance simulator. https://engineering.fb.com/2023/09/07/ data-infrastructure/arcadia-end-to-end-ai-system-performancesimulator/ Meta Engineering Blog Accessed: 2026-02-06. [17] Yicheng Feng, Yuetao Chen, Kaiwen Chen, Jingzong Li, Tianyuan Wu, Peng Cheng, Chuan Wu, Wei Wang, Tsung-Yi Ho, and Hong Xu. 2024. Echo: Simulating Distributed Training At Scale. arXiv preprint arXiv:2412.12487 (2024). [18] Adithya Gangidi, Rui Miao, Shengbao Zheng, Sai Jayesh Bondu, Guilherme Goes, Hany Morsy, Rohit Puri, Mohammad Riftadi,

Ashmitha Jeevaraj Shetty, Jingyi Yang, et al. 2024. Rdma over ethernet for distributed training at meta scale. In Proceedings of the ACM SIGCOMM 2024 Conference. 57–70. [19] Amir Gholami, Zhewei Yao, Sehoon Kim, Coleman Hooper, Michael W Mahoney, and Kurt Keutzer. 2024. AI and memory wall. IEEE Micro (2024). [20] Alibaba Group. aliyun/aicb at d9b4f5cd7d9d34a80cfbb0389831a16c7fe3ed7b. https://github.com/ aliyun/aicb/tree/d9b4f5cd7d9d34a80cfbb0389831a16c7fe3ed7b Accessed: 2026-02-06. [21] 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 22nd USENIX Symposium on Networked Systems Design and Implementation (NSDI 25). 473–488. [22] Runsheng Benson Guo, Utkarsh Anand, Arthur Chen, and Khuzaima Daudjee. 2024. Cephalo: Harnessing Heterogeneous GPU Clusters for Training Transformer Models. arXiv preprint arXiv:2411.01075 (2024). [23] Tom’s Hardware. 2025. Microsoft deploys world’s first “supercomputerscale” GB300 NVL72 Azure cluster — 4,608 GB300 GPUs linked together to form a single, unified accelerator capable of 1.44 PFLOPS of inference. https://www.tomshardware.com/tech-industry/artificialintelligence/microsoft-deploys-worlds-first-supercomputer-scalegb300-nvl72-azure-cluster-4-608-gb300-gpus-linked-togetherto-form-a-single-unified-accelerator-capable-of-1-44-pflops-ofinference Accessed: 2026-02-06. [24] Aaron Harlap, Deepak Narayanan, Amar Phanishayee, Vivek Seshadri, Nikhil Devanur, Greg Ganger, and Phil Gibbons. 2018. Pipedream: Fast and efficient pipeline parallel dnn training. arXiv preprint arXiv:1806.03377 (2018). [25] Chenyang Hei, Jiayi Li, Jiamin Cao, Chengxi Gao, Xiuzhu Sha, Tongrui Liu, Dengke Zhang, Ennan Zhai, and Xingwei Wang. 2026. HeteCCL: Synthesizing Near-Optimal Collective Communication Schedules for Heterogeneous GPU Clusters. In 23rd USENIX Symposium on Networked Systems Design and Implementation (NSDI 26). USENIX Association, Renton, WA, 2533–2551. https://www.usenix.org/conference/ nsdi26/presentation/hei [26] Myeongjae Jeon, Shivaram Venkataraman, Amar Phanishayee, Junjie Qian, Wencong Xiao, and Fan Yang. 2019. Analysis of Large-Scale Multi-Tenant GPU Clusters for DNN Training Workloads. In 2019 USENIX Annual Technical Conference (USENIX ATC 19). USENIX Association, Renton, WA, 947–960. https://www.usenix.org/conference/ atc19/presentation/jeon [27] Xianyan Jia, Le Jiang, Ang Wang, Wencong Xiao, Ziji Shi, Jie Zhang, Xinyuan Li, Langshi Chen, Yong Li, Zhen Zheng, et al. 2022. Whale: Efficient giant model training over heterogeneous {GPUs}. In 2022 USENIX Annual Technical Conference (USENIX ATC 22). 673–688. [28] Stefan Näher Kurt Mehlhorn. 1994. Implementation of a Sweep Line Algorithm for the Straight Line Segment Intersection Problem. MAXPLANCK-INSTITUT•• FUR INFORMATIK. https://pure.mpg.de/rest/ items/item_1834220_3/component/file_2035159/content [29] Seonho Lee, Amar Phanishayee, and Divya Mahajan. 2025. Forecasting GPU performance for deep learning training and inference. In Proceedings of the 30th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 1. 493–508. [30] Wenkai Li, Ran Shu, Peng Zhang, and Yongqiang Xiong. 2025. Nüwa: Efficient Generative Control Plane for AI Network Simulation. In Proceedings of the 9th Asia-Pacific Workshop on Networking. 121–127. [31] Ji Liu, Zhihua Wu, Danlei Feng, Minxu Zhang, Xinxuan Wu, Xuefeng Yao, Dianhai Yu, Yanjun Ma, Feng Zhao, and Dejing Dou. 2023. Heterps: Distributed deep learning with reinforcement learning based scheduling in heterogeneous environments. Future Generation Computer Systems 148 (2023), 106–117. 13

[32] Fei Long, Kaihui Gao, Li Chen, Dan Li, Yiwei Zhang, Fei Gui, Yitao Xing, Wenjia Wei, and Bingyang Liu. 2026. Supercharging Packet-level Network Simulation of Large Model Training via Memoization and Fast-Forwarding. arXiv preprint arXiv:2602.10615 (2026). [33] Yixuan Mei, Yonghao Zhuang, Xupeng Miao, Juncheng Yang, Zhihao Jia, and Rashmi Vinayak. 2025. Helix: Serving Large Language Models over Heterogeneous GPUs and Network via Max-Flow. In Proceedings of the 30th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 1. 586–602. Meta’s Infrastructure Evolution and the Advent of [34] Meta. AI. https://engineering.fb.com/2025/09/29/data-infrastructure/metasinfrastructure-evolution-and-the-advent-of-ai/. Accessed: 2026-0206. [35] Zizhao Mo, Huanle Xu, and Chengzhong Xu. 2024. Heet: Accelerating elastic training in heterogeneous deep learning clusters. In Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2. 499–513. [36] Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Korthikanti, Dmitri Vainbrand, and Bryan Catanzaro. 2021. Scaling Language Model Training to a Trillion Parameters Using Megatron. https://developer.nvidia.com/blog/scalinglanguage-model-training-to-a-trillion-parameters-using-megatron/ Accessed: 2026-02-06. [37] Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, et al. 2021. Efficient large-scale language model training on gpu clusters using megatronlm. In Proceedings of the international conference for high performance computing, networking, storage and analysis. 1–15. [38] NVIDIA Networking. RDMA over Converged Ethernet (RoCE) v2. https://docs.nvidia.com/doca/archive/2-10-0/rdma-over-convergedethernet/index.html Accessed: 2026-02-06. [39] Chengyi Nie, Jessica Maghakian, and Zhenhua Liu. 2024. Cannikin: Optimal Adaptive Distributed DNN Training over Heterogeneous Clusters. In Proceedings of the 25th International Middleware Conference. 299–312. [40] NVIDIA. DGX SuperPOD Reference Architecture: DGX H100. https://docs.nvidia.com/https:/docs.nvidia.com/dgx-superpodreference-architecture-dgx-h100.pdf. Accessed: 2026-02-06. [41] Nvidia. NVIDIA A100 TENSOR CORE GPU. https://www.nvidia.com/content/dam/en-zz/Solutions/DataCenter/a100/pdf/nvidia-a100-datasheet-us-nvidia-1758950-r4web.pdf. Accessed: 2026-02-06. [42] NVIDIA. NVIDIA H100 Tensor Core GPU. https://www.nvidia.com/enin/data-center/h100/. Accessed: 2026-02-06. [43] Jay H Park, Gyeongchan Yun, M Yi Chang, Nguyen T Nguyen, Seungmin Lee, Jaesik Choi, Sam H Noh, and Young-ri Choi. 2020. {HetPipe}: Enabling large {DNN} training on (whimpy) heterogeneous {GPU} clusters through integration of pipelined model parallelism and data parallelism. In 2020 USENIX Annual Technical Conference (USENIX ATC 20). 307–321. [44] Guicheng Qi, Junwei Su, Liqi Yang, Tao Li, Tingwen Xie, Yerui Sun, Yuchen Xie, and Chuan Wu. 2026. HetAuto: Cross-Cluster AutoParallelism for Heterogeneous Distributed Training. In Proceedings of the 21st European Conference on Computer Systems. 759–779. [45] Yicheng Qian, Ran Shu, Rui Ma, Yang Wang, Derek Chiou, Nadeen Gebara, Luca Piccolboni, Miriam Leeser, and Yongqiang Xiong. 2025. Miniature: Fast AI Supercomputer Networks Simulation on FPGAs. In Proceedings of the 9th Asia-Pacific Workshop on Networking. 114–120. [46] Jianxing Qin, Jingrong Chen, Xinhao Kong, Yongji Wu, Tianjun Yuan, Liang Luo, Zhaodong Wang, Ying Zhang, Tingjun Chen, Alvin R Lebeck, et al. 2025. Phantora: Maximizing Code Reuse in Simulationbased Machine Learning System Performance Estimation. arXiv preprint arXiv:2505.01616 (2025).

[47] 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. 3505–3506. [48] sagemaker. sagemaker. https://aws.amazon.com/blogs/machinelearning/improve-price-performance-of-your-model-trainingusing-amazon-sagemaker-heterogeneous-clusters/. Accessed: 2026-02-06. [49] Amazon Web Services. 2025. Amazon Virtual Private Cloud (VPC) Overview. https://docs.aws.amazon.com/vpc/latest/userguide/whatis-amazon-vpc.html Accessed: 2026-02-06. [50] Siyuan Shen, Tommaso Bonato, Zhiyi Hu, Pasquale Jordan, Tiancheng Chen, and Torsten Hoefler. 2025. Atlahs: An application-centric network simulator toolchain for ai, hpc, and distributed storage. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis. 349–367. [51] 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). [52] Min Si, Pavan Balaji, Yongzhou Chen, Ching-Hsiang Chu, Adi Gangidi, Saif Hasan, Subodh Iyengar, Dan Johnson, Bingzhe Liu, Regina Ren, Ashmitha Jeevaraj Shetty, Greg Steinbrecher, Yulun Wang, Bruce Wu, Xinfeng Xie, Jingyi Yang, Mingran Yang, Kenny Yu, Minlan Yu, Cen Zhao, Wes Bland, Denis Boyda, Suman Gumudavelli, Prashanth Kannan, Cristian Lumezanu, Rui Miao, Zhe Qu, Venkat Ramesh, Maxim Samoylov, Jan Seidel, Srikanth Sundaresan, Feng Tian, Qiye Tan, Shuqiang Zhang, Yimeng Zhao, Shengbao Zheng, Art Zhu, and Hongyi Zeng. 2025. Collective Communication for 100k+ GPUs. arXiv:2510.20171 [cs.DC] https://arxiv.org/abs/2510.20171 [53] Arjun Singhvi, Nandita Dukkipati, Prashant Chandra, Hassan MG Wassel, Naveen Kr Sharma, Anthony Rebello, Henry Schuh, Praveen Kumar, Behnam Montazeri, Neelesh Bansod, et al. 2025. Falcon: A reliable, low latency hardware transport. In Proceedings of the ACM SIGCOMM 2025 Conference. 248–263. [54] 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). [55] Foteini Strati, Zhendong Zhang, George Manos, Ixeia Sánchez Périz, Qinghao Hu, Tiancheng Chen, Berk Buzcu, Song Han, Pamela Delgado, and Ana Klimovic. 2025. Sailor: Automating distributed training over dynamic, heterogeneous, and geo-distributed clusters. In Proceedings of the ACM SIGOPS 31st Symposium on Operating Systems Principles. 204–220. [56] Yinan Tang, Tongtong Yuan, Fang Cao, Li Wang, Zhenhua Guo, Yaqian Zhao, and Rengang Li. 2024. Simulating llm training in cxl-based heterogeneous computing cluster. In IEEE INFOCOM 2024-IEEE Conference on Computer Communications Workshops (INFOCOM WKSHPS). IEEE, 1–6. [57] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. 2023. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 (2023). [58] Taegeon Um, Byungsoo Oh, Minyoung Kang, Woo-Yeon Lee, Goeun Kim, Dongseob Kim, Youngtaek Kim, Mohd Muzzammil, and Myeongjae Jeon. 2024. Metis: Fast Automatic Distributed Training on Heterogeneous {GPUs}. In 2024 USENIX Annual Technical Conference (USENIX ATC 24). 563–578. [59] Wenyi Wang, Zheng Wu, Yanmeng Wang, Haolin Mao, Lei Han, Gaogang Xie, and Fu Xiao. 2026. HyGra: Accelerating Network-State 14

[73] Liangyu Zhao, Saeed Maleki, Yuanhong Wang, Zezhou Wang, Ziyue Yang, Hossein Pourreza, and Arvind Krishnamurthy. 2024. Forestcoll: throughput-optimal collective communications on heterogeneous network fabrics. arXiv preprint arXiv:2402.06787 (2024). [74] Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, et al. 2023. Pytorch fsdp: experiences on scaling fully sharded data parallel. arXiv preprint arXiv:2304.11277 (2023). [75] Yonghao Zhuang, Lianmin Zheng, Zhuohan Li, Eric Xing, Qirong Ho, Joseph Gonzalez, Ion Stoica, Hao Zhang, and Hexu Zhao. 2023. On optimizing the communication of model parallelism. Proceedings of Machine Learning and Systems 5 (2023), 526–540.

Simulation for LLM Training in DCNs via Adaptive Packet-Flow Granularity. arXiv preprint arXiv:2603.12671 (2026). [60] Xizheng Wang, Qingxu Li, Yichi Xu, Gang Lu, Dan Li, Li Chen, Heyang Zhou, Linkang Zheng, Sen Zhang, Yikai Zhu, et al. 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). 541–558. [61] Qizhen Weng, Wencong Xiao, Yinghao Yu, Wei Wang, Cheng Wang, Jian He, Yong Li, Liping Zhang, Wei Lin, and Yu Ding. 2022. MLaaS in the Wild: Workload Analysis and Scheduling in Large-Scale Heterogeneous GPU Clusters. In 19th USENIX Symposium on Networked Systems Design and Implementation (NSDI 22). USENIX Association, Renton, WA, 945–960. https://www.usenix.org/conference/nsdi22/ presentation/weng [62] 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). IEEE, 283–294. [63] Ruilong Wu, Xinjiao Li, Yisu Wang, Xinyu Chen, and Dirk Kutscher. 2025. Rethinking Dynamic Networks and Heterogeneous Computing with Automatic Parallelization. In Proceedings of the 9th Asia-Pacific Workshop on Networking. 164–171. [64] Yongji Wu, Xueshen Liu, Shuowei Jin, Ceyu Xu, Feng Qian, Z Morley Mao, Matthew Lentz, Danyang Zhuo, and Ion Stoica. 2025. HeterMoE: Efficient Training of Mixture-of-Experts Models on Heterogeneous GPUs. arXiv preprint arXiv:2504.03871 (2025). [65] Ran Yan, Youhe Jiang, Xiaonan Nie, Fangcheng Fu, Bin Cui, and Binhang Yuan. 2024. HexiScale: Accommodating Large Language Model Training over Heterogeneous Environment. arXiv preprint arXiv:2409.01143 (2024). [66] Ran Yan, Youhe Jiang, Wangcheng Tao, Xiaonan Nie, Bin Cui, and Binhang Yuan. 2024. Flashflex: Accommodating large language model training over heterogeneous environment. arXiv e-prints (2024), arXiv– 2409. [67] Shengyuan Ye, Liekang Zeng, Xiaowen Chu, Guoliang Xing, and Xu Chen. 2024. Asteroid: Resource-efficient hybrid pipeline parallelism for collaborative DNN training on heterogeneous edge devices. In Proceedings of the 30th Annual International Conference on Mobile Computing and Networking. 312–326. [68] Xiaodong Yi, Shiwei Zhang, Ziyue Luo, Guoping Long, Lansong Diao, Chuan Wu, Zhen Zheng, Jun Yang, and Wei Lin. 2020. Optimizing distributed training deployment in heterogeneous GPU clusters. In Proceedings of the 16th International Conference on emerging Networking EXperiments and Technologies. 93–107. [69] Xiaofei Yue, Fangming Zhao, Fulun Ye, Jiongchi Yu, Zhaoxuan Li, Tingting Li, Ziming Zhao, and Jianwei Yin. 2026. HeteroSim: Towards High-Fidelity Heterogeneous LLM Training Simulation on GPUs. In Proceedings of the ACM Web Conference 2026. 5189–5197. [70] Jinghui Zhang, Geng Niu, Qiangsheng Dai, Haorui Li, Zhihua Wu, Fang Dong, and Zhiang Wu. 2023. PipePar: Enabling fast DNN pipeline parallel training in heterogeneous GPU clusters. Neurocomputing 555 (2023), 126661. [71] Shiwei Zhang, Lansong Diao, Chuan Wu, Zongyan Cao, Siyu Wang, and Wei Lin. 2024. Hap: Spmd dnn training on heterogeneous gpu clusters with automated program synthesis. In Proceedings of the Nineteenth European Conference on Computer Systems. 524–541. [72] Zili Zhang, Yinmin Zhong, Yimin Jiang, Hanpeng Hu, Jianjian Sun, Zheng Ge, Yibo Zhu, Daxin Jiang, and Xin Jin. 2025. DistTrain: Addressing Model and Data Heterogeneity with Disaggregated Training for Multimodal Large Language Models. In Proceedings of the ACM SIGCOMM 2025 Conference. 24–38.

A

Input Specification Sample

The simulation framework mainly works with three input specifications which define our heterogeneous deployment configurations. Each configuration file captures distinct training setup, which enables us to give the simulator information about hybrid gpu clusters. Step 1. How are device groups and parallelism metadata defined? The first step initiates with creating a protocol buffer specification file that defines the device group configurations. A device group represents a collection of GPUs among which tensor parallelism must happen. As shown in the Figure 13, each device group entry specifies the global GPU ranks belonging to the group (e.g, ranks {0, 1, 2} form a tensor parallel group of size 3) along with parallelism metadata. For instance, our given configuration with ranks {0, 1, 2} at pipeline stage 1 and replica 1, execute the first set of transformer layers. This structured approach helps AICB to generate heterogeneous device group specific workloads. Step 2. Next! How can we model the heterogeneous network structure? This step involves creation of the topology file providing comprehensive network information for the simulation. From the Figure 13 the file header specifies global topology parameters: total node count which consists of GPUs, NVLinks, and switches. Following the header, NVLink node identifier nodes are listed, then individual connection entries define the network structure. Adding our implementation the topology file captures heterogeneous link parameters reflecting real hardware differences, for example, H100 GPUs with 3200 Gbps NVLink bandwidth and 0.000025ms latency versus A100 GPUs with 256 Gbps NVLink bandwidth and 0.0005ms latency. Step 3. Finally! How are specific workloads generated for each group? The third steps the the model specifications and parameters which we provide AICB with to generate per device group workload files. Model parameters include hidden size (e.g., 4096 for LLaMA-7B), FFN hidden size (16384), number of layers (32), attention heads (32), vocabulary size (32000), and sequence length (2048). The workload generator which is built upon AICB (AI Collective Bechmark), 15

Figure 13. Input Specification for the example deployment configuration. yields: Punique = {1, 16, 21, 33}. Let 𝑝𝑖 be the 𝑖 th boundary point in Punique . From these boundary points, we induce the contiguous layer segments: [𝑝 0, 𝑝 1 − 1] = {1, 15} [𝑝 1, 𝑝 2 − 1] = {16, 20} [𝑝 2, 𝑝 3 − 1] = {21, 32} Step 3: Cool! Now which Device Groups cover which segments? 𝐹𝑜𝑟 𝑒𝑎𝑐ℎ segment [𝑠𝑒𝑔start, 𝑠𝑒𝑔end ], Algorithm 1 constructs a candidate set 𝐶 consisting of all the Device Groups whose layer ranges fully cover the segment. Segment [1,15]: Both 𝐷𝐺 0 [1, 20] and 𝐷𝐺 2 [1, 15] fully cover this segment, so: 𝐶 = {𝐷𝐺 0, 𝐷𝐺 2 }. Segment [16,20]: This segment is fully covered by 𝐷𝐺 0 [1, 20] and 𝐷𝐺 3 [16, 32], so: 𝐶 = {𝐷𝐺 0, 𝐷𝐺 3 }. Segment [21,32]: This segment is covered by 𝐷𝐺 1 [21, 32] and 𝐷𝐺 3 [16, 32], so: 𝐶 = {𝐷𝐺 1, 𝐷𝐺 3 }.

processes this parameters along with the device group configurations to generate workload file per device group. Each of the workload file contains the training iteration structure: a header specifying the tensor parallel degree, pipeline parallel stages, gradient accumulation steps and embedding parameters. The header is then followed by work items for each layer including attention layers, MLP layers, and optimizer steps. Each work item encode with compute time(via GPU profiling), communication type (ALLREDUCE, ALLGATHER, REDUCESCATTER), communication size in bytes.

B

Heterogeneity-aware Multi-Ring Construction Example

Consider the deployment (Fig. 14) of a model with 32 layers across four heterogeneous Device Groups with the following assigned layer ranges [𝑠𝑖 , 𝑒𝑖 ]: 𝐷𝐺 0 with layer range [1, 20] and ranks {0, 1, 2}, 𝐷𝐺 1 with layer range [21, 32] and ranks {3, 4}, 𝐷𝐺 2 with layer range [1, 15] and ranks {5, 6} and, 𝐷𝐺 3 with layer range [16, 32] and ranks {7, 8, 9}. This represents an asymmetric pipeline partitioning deployment and therefore, these Device Groups overlap on different portions of the model, implying that different subsets of ranks must synchronize gradients for different layer ranges. The sweep-line algorithm (Algorithm 1) essentially answers the question "How do we find all these subsets?" We illustrate the sweep-line algorithm using the above example.

Step 3: Nice! Let’s finally construct our DP synchronization groups 𝐹𝑜𝑟 𝑒𝑎𝑐ℎ segment, if |C| ≥ 2, Algorithm 1 constructs a DP synchronization group G𝐷𝑃 by taking the union of ranks across all Device Groups in 𝐶. Segment [1,15]: G𝐷𝑃 = {0, 1, 2} ∪ {3, 4} = {0, 1, 2, 3, 4} Segment [16,20]: G𝐷𝑃 = {0, 1, 2} ∪ {7, 8, 9} = {0, 1, 2, 7, 8, 9} Segment [21,32]: G𝐷𝑃 = {5, 6} ∪ {7, 8, 9} = {5, 6, 7, 8, 9} Voila! There we have our DP groups. Each G𝐷𝑃 is associated with its corresponding layer segment.

Step 1: Can we arrange the layer boundary points on a number line? Algorithm 1 begins by collecting the layer boundary points from all Device Groups. 𝐹𝑜𝑟 𝑒𝑎𝑐ℎ Device Group with range [𝑠𝑖 , 𝑒𝑖 ], the algorithm inserts 𝑠𝑖 and (𝑒𝑖 + 1) into the boundary set P. We increment the end boundary points (𝑒𝑖 ) by one which ensures that adjacent segments are handled cleanly. For this example, P = {1, 21, 21, 33, 1, 16, 16, 33}. Step 2: Looks a bit messy. Let’s sort and clean! The boundary set P is then sorted and de-duplicated, which

We illustrate Algorithm 2 using the following. Consider a DP synchronization group which is responsible for the layer range [1, 15]. This group consists of two Device Groups: 𝐷𝐺 0 with tensor parallelism degree TP=3 and ranks {0, 1, 2} and 𝐷𝐺 2 with 𝑇 𝑃 = 2 and ranks {3, 4}.

16

Figure 14. Heterogeneity-aware multi-ring LCM-based resharding in 𝑋𝑠𝑖𝑚 for non-uniform layer partitioning across highcompute (blue) and low-compute (red) device groups with layer-aware DP groups. With the above mapping, Algorithm 2 constructs one communication ring per chunk. Each ring connects exactly one rank from 𝐷𝐺 0 and one rank from 𝐷𝐺 2 that owns the same chunk, ensuring that all participants in a ring operate on identically sized tensors. The interleaved assignment helps in distributing all communication evenly across the ranks and assign multiple chunks to a single rank.

Step 1: Which Device Groups participate? For the layer range [1, 15], the sweep-line algorithm has already determined that 𝐷𝐺 0 and 𝐷𝐺 2 will be synchronizing gradients. Algorithm 2 therefore takes these 2 Device Groups as the input DP group. Step 2: Wait! But how many chunks and rings? Since the participating Device Groups use different tensor parallelism degrees, their gradient tensors are partitioned differently. Algorithm 2 resolves this mismatch by computing the least common multiple (𝐿𝐶𝑀) of the TP degrees of the participating Device Groups. In this example, 𝑙𝑐𝑚(3, 2) = 6. This value defines the number of gradient chunks to be communicated across the participating Device Groups inside the DP group and, hence, the number of communication rings required for the correct synchronization. Algorithm 2 conceptually represents each gradient as 6 equal sized chunks. Ranks in 𝐷𝐺 0 collectively cover these chunks with 2 chunks per rank, and ranks in 𝐷𝐺 2 cover 6 chunks with 3 chunks per rank. This representation helps in abstracting the physical tensor layouts while ensuring compatibility in synchronization across Device Groups. Step 3: Okay! Now, which chunk goes to which rank? Algorithm 2 assigns chunks to ranks in an interleaved/roundrobin fashion based on each rank’s local index within its Device Group. In this example, chunks are indexed from 0 to 5, and for each chunk 𝑐, the algorithm selects the participating ranks in each Device Group as local_rank = 𝑐 mod TP. For 𝐷𝐺 0 , this maps chunks {0, 3} to local rank 0, {1, 4} to local rank 1, and {2, 5} to local rank 2. For 𝐷𝐺 2 , chunks {0, 2, 4} map to local rank 0 and chunks {1, 3, 5} to local rank 1. Step 4: Where are my rings?

C

Heterogeneity-aware Chunk Partitioning Example

We illustrate Algorithm 3 using the example of the DP synchronization group above. Consider a DP synchronization group GD P responsible for synchronizing gradients for the layer range [1, 50]. As per the sweep-line algorithm (Algorithm 1), this group consists of two Device Groups: 𝐷𝐺 0 : TP degree 𝑡 0 = 3, ranks {0, 1, 2} 𝐷𝐺 2 : TP degree 𝑡 1 = 3, ranks {3, 4} Let the total gradient communication volume for this DP group be 𝑑 = 60 𝑀𝐵. Since the two Device Groups have different TP degrees, their gradients are initially partitioned into incompatible tensor shapes and hence, cannot be directly aggregated using a standard collective like AllReduce. The goal of Algorithm 3 is to resolve this mismatch into a common number of finer grained equal sized chunks. Step 1: How many chunks should be there in my GD P ? Algorithm 3 starts by computing the least common multiple 17

of the participating TP degrees: 𝐿 = 𝑙𝑐𝑚(𝑡 0, 𝑡 1 ) = 𝑙𝑐𝑚(3, 2) = 6. 𝐿 represents the total number of local gradient chunks used for synchronization with respect to GD P . This ensures tensor compatibility across heterogeneous TP configurations. Step 2: Umm...What is my per-rank gradient volume per Device Group? 𝐹𝑜𝑟 𝑒𝑎𝑐ℎ Device Group 𝐷𝐺𝑖 ∈ GD P , Algorithm 3 computes the gradient volume handled by a single rank in 𝐷𝐺𝑖 :

matches SimAI across all configurations, with relative errors between 0.1–2.2%, validating the correctness of 𝑋𝑠𝑖𝑚 ’s end-to-end modeling pipeline for homogeneous clusters. (8.b) Comparing 𝑋𝑠𝑖𝑚’s simulation time with SOTA training simulator, SimAI, for homogeneous AI clusters. Fig. 16 demonstrates the wall-clock time of the simulator, i.e., the simulation runtime for Llama 7B over varying cluster sizes. At large cluster scales (i.e., 128 and 512 GPUs), 𝑋𝑠𝑖𝑚 achieves lower per-iteration simulation runtime than SimAI. This improvement stems from 𝑋𝑠𝑖𝑚 ’s LCM-based chunking and chunk partitioning, which limit chunk fragmentation and reduce the number of simulated communication events. In contrast, SimAI’s mockNCCL implementation increases the number of chunks as the cluster size grows, even for small models, leading to higher eventscheduling and synchronization overhead during simulation. As a result, 𝑋𝑠𝑖𝑚 incurs lower wall-clock simulation time at scale. These differences are negligible at small cluster sizes but become pronounced as collective communication dominates execution at higher degrees of parallelism. (9) 𝑋𝑠𝑖𝑚’s scalability analysis. We evaluate 𝑋𝑠𝑖𝑚’s scalability across model sizes (Llama 7B and Llama 13B) and heterogeneous cluster configurations, i.e., 𝐶15, with varying cluster size (see Fig. 17). As cluster size increases, the simulation runtime decreases substantially, achieving up to 5.75× speedup for Llama 7B and 6.25× for Llama 13B models. For any fixed cluster, the Llama 13B model incurs higher runtime due to increased computation and communication. The performance gap narrows at mid-scale configurations due to fewer simulated network events per iteration. As cluster size increases, collective communication is partitioned across more ranks and rings, resulting in smaller per-flow data sizes and more regular traffic patterns. These results demonstrate that 𝑋𝑠𝑖𝑚 scales efficiently with both model and cluster size in heterogeneous environments. (10.a) 𝑋𝑠𝑖𝑚’s utility: How can GPU idle time or straggler analysis help deployers? Fig. 18 shows the model training time, and the GPU idle time (i.e., the Straggler waiting time) across three heterogeneous configurations, 𝐶13 (DP), 𝐶14 (TP+DP), and 𝐶15 (TP+PP) (see Table 4). We observe that the heterogeneous DP configuration exhibits a waiting time of 4.42 𝑠𝑒𝑐, while switching to a TP+DP configuration reduces waiting time to 0.294 𝑠𝑒𝑐, and further reduces to 0.07 𝑠𝑒𝑐 under PP+TP, demonstrating how workload partitioning mitigates synchronization stalls. More interestingly, we see that the training time for 𝐶13 and 𝐶14 is similar, but the GPUs with 𝐶13 are underutilized. One may optimize the algorithm or use GPU time for alternative tasks, such as checkpointing. (10.b) 𝑋𝑠𝑖𝑚’s utility: Total Cost of Ownership (TCO) Across homogeneous and heterogeneous cluster designs using Llama 7B (see Fig. 19), we observe that heterogeneous clusters can achieve near-identical or superior performance

𝑑 60 𝑀𝐵 = 20 𝑀𝐵 = 𝑡0 3 60 𝑀𝐵 𝑑 = 30 𝑀𝐵 For 𝐷𝐺 2 : data_per_rank𝐷𝐺 2 = = 𝑡2 2 For 𝐷𝐺 0 : data_per_rank𝐷𝐺 0 =

This ensures that under tensor parallelism, each rank initially owns an equal fraction of the gradient within its Device Group. Step 3: But (data_per_rank𝐷𝐺 0 ≠ data_per_rank𝐷𝐺 2 ). So how do we make the chunk sizes equal? Next, Algorithm 3 determines how many chunks each rank in 𝐿 a Device Group will participate in: chunk_multiplier𝐷𝐺𝑖 = 𝑡𝑖 𝐿 6 For 𝐷𝐺 0 : chunk_multiplier𝐷𝐺 0 = = = 3 𝑡0 2 𝐿 6 For 𝐷𝐺 2 : chunk_multiplier𝐷𝐺 2 = = = 2 𝑡2 3 This means that each rank in 𝐷𝐺 0 contributes to 2 chunks and each rank in 𝐷𝐺 2 contributes to 3 ranks. From this information, Algorithm 3 computes the volume of data contributed by each rank to each chunk: data_per_chunk_per_rank =

data_per_rank𝐷𝐺𝑖

. chunk_multiplier𝐷𝐺𝑖 20 𝑀𝐵 = 10 𝑀𝐵 For 𝐷𝐺 0 : data_per_chunk_per_rank = 2 30 𝑀𝐵 For 𝐷𝐺 2 : data_per_chunk_per_rank = = 10 𝑀𝐵 3 Thus, regardless of the uneven TP degrees, every participating rank contributes exactly 10 MB to each chunk it owns.

D

Additional Evaluation Results

(8.a) Comparing 𝑋𝑠𝑖𝑚’s prediction accuracy with SOTA training simulator, SimAI, for homogeneous AI clusters. Given that SimAI [60] has demonstrated high predictive accuracy for homogeneous clusters using real-world experiments, we use it as a reference to validate the correctness of 𝑋𝑠𝑖𝑚 ’s modeling pipeline. Fig. 15 reports the predicted training time (per iteration) for Llama 7B and Llama 13B across varying homogeneous cluster sizes. 𝑋𝑠𝑖𝑚 closely 18

Llama 13B

2.19% 0.30%

6000 4000

0.41%

0.45%

2000

Llama 7B

1.97%

0

8

16 128 Cluster Size

512

Figure 15. Training time per iteration for Llama 7B and 13B across homogeneous cluster sizes shows that 𝑋𝑠𝑖𝑚 closely matches SimAI, with a relative error of 0.1–2.2%.

5 Llama 7B

SimAI

Xsim

Llama 13B

Simulation runtime (per iteration) (hour)

0.76%

Simulation runtime (per iteration) (hour)

0.11%

Xsim

GBS = 128 GBS = 2048

Training time per iteration (ms)

1.80%

GBS = 128 GBS = 2048

SimAI

12500 10000 7500 5000 2500 0

4

4

3

3

2

2

8

16 128 Cluster size

0

512

Figure 16. Simulation runtime per iteration for Llama 7B across homogeneous cluster sizes shows that 𝑋𝑠𝑖𝑚 achieves lower runtime at large scales (128–512 GPUs) due to LCM-based chunking and reduced communication overhead, with negligible differences at small scales.

256

1.120.93 512

Cluster size

0.68 0.36 1024

ONLY A100

A100 & H100

C5 C6 C14

C7 C8 C16

300000

$/GPU-hour

6000

Time Training Time GPU Idle (ms) per Iteration (sec)

128

1.85 1.25

Figure 17. Simulation runtime per iteration versus heterogeneous cluster size for Llama 7B and 13B shows that 𝑋𝑠𝑖𝑚 scales efficiently, achieving up to 5.75× speedup for 7B and 6.25× for 13B as cluster size increases.

ONLY H100 200000

4000

100000

2000

0 80

Training time per iteration (sec)

0 100 75 50 25 0

2.07

1

1 0

Llama 7B

4.27

60 40 20

C13

C14 Cluster Configuration

C15

0

C3 C4 C13

Cluster Configuration

Figure 18. GPU idle time and training time per iteration across cluster configurations ((C13–C15) Table 4), highlighting imbalance and utilization differences.

Figure 19. Cost per GPU-hour and training time per iteration for Llama 7B across homogeneous and heterogeneous clusters.

while significantly reducing cost. For instance, in the 8-GPU data-parallel configuration (𝐶3, 𝐶4, 𝐶13), a mixed 4×H100 + 4×A100 cluster achieves a completion time of 57.48 s, comparable to the 8×H100 cluster (57.06 s), yet reduces TCO from 210𝐾 $/GPU-hour to 146𝐾 $/GPU-hour, yielding a ∼30% costefficiency improvement. Similarly, compared to an 8×A100 cluster (TCO = 83𝐾 $/GPU-hour), the heterogeneous configuration offers higher performance headroom at moderate additional cost. Interestingly, for the configurations 𝐶9, 𝐶10, and 𝐶16, we observe that the heterogeneity solution performs best and has a lower cost than the 8×H100 cluster. This could be because the non-uniform hybrid parallelism led to better resource utilization.

E 𝑋𝑠𝑖𝑚’s LCM-based tensor resharding: Bounds and Tradeoffs Since practical deployments use TP degree, 𝑡𝑖 ≤ 8, the maximum possible synchronization granularity is bounded by LCM(𝑡 1, . . . , 𝑡𝑘 ) ≤ 8 × 7 × 5 × 3 = 840, which occurs when the TP degrees collectively contain all prime factors up to 8. Thus, the worst-case LCM remains bounded by 840 regardless of the number of participating device groups (𝑘 ≥ 4). For an AllReduce with 𝑘 participants and message size 𝑐, the communication cost is approximately
𝑐 𝑇ring ≈ 2(𝑘 − 1) 𝛼 + 𝑘𝐵 , for a ring algorithm, and
𝑇tree ≈ 2 log2 (𝑘) 𝛼 + 𝐵𝑐 , for a binary-tree algorithm, where 𝛼 is the per-message latency and 𝐵 is the link bandwidth. In large-scale deployments (> 16𝐾 GPUs), hierarchical collectives typically use local groups with 𝑡𝑖 ≤ 8, resulting 19

in a global DP group of roughly 16,000/8 ≈ 2,000 participants. Further sharding (e.g., FSDP [74]) commonly operates on smaller subgroups (e.g., 128 ranks). For a 4 GB gradient synchronization (approximately the gradient size of Llama-2 GB ≈ 32 MB, indicat70B), the per-rank chunk size is only 4128 ing that the bounded LCM and chunk-partitioning overheads remain practical even at large scales.

F

Extended Related Work

High-speed simulators. Nüwa [30] accelerates AI network simulation by reducing routing construction overhead via formula-based routing. Multiverse [21] achieves speedups using a GPU-accelerated design. FPGA-based simulators such as Miniature [45] achieve orders-of-magnitude acceleration by mapping homogeneous network components to specialized hardware, scaling to 65K nodes. Recent packet-level simulators such as Wormhole [32] and HyGra [59] accelerate LLM network simulation through steady-state detection, state reuse, network partitioning, and adaptive switching between packet- and flow-level models. While these techniques significantly reduce simulation time, they assume homogeneous infrastructure and fixed communication semantics. 𝑋𝑠𝑖𝑚 complements these efforts by providing heterogeneityaware full-stack simulation, while supporting htsim as a scalable software backend for faster network simulation without specialized hardware, albeit with lower acceleration than hardware-assisted approaches.

20