Simulating Unified Tensor Resharding in heterogeneous AI systems

2026-06-25

AI Summary

This work addresses the limitation of existing AI training simulators, which often assume homogeneous compute and network infrastructure and thus fail to accurately model real-world heterogeneous training environments. To overcome this, the authors propose Xsim—the first distributed simulator supporting heterogeneous large model training—that innovatively unifies the simulation of non-uniform tensor resharding and pipeline parallelism. Xsim introduces heterogeneity-aware communication modeling, reusable parallel algorithm abstractions, and a flexible deployment mechanism. It integrates techniques such as custom ring construction, block partitioning, and plugins from NS-3 and htsim to achieve high-fidelity, scalable performance simulation. Experimental results demonstrate that Xsim achieves training time prediction errors below 5% across diverse heterogeneous configurations, with pipeline communication modeling errors around 2%, while also providing key metrics such as pipeline bubble duration and straggler waiting time.

Abstract

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 Xsim, a heterogeneity-aware simulator for distributed LLM training. Xsim 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 Xsim 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.

Problem

• heterogeneous AI systems
• distributed LLM training
• tensor resharding
• training simulation
• infrastructure heterogeneity

Innovation

• heterogeneity-aware simulation
• non-uniform tensor resharding
• distributed LLM training
• customized collective communication
• pipeline-parallel algorithms