# KAITIAN: A Unified Communication Framework for Enabling Efficient Collaboration Across Heterogeneous Accelerators in Embodied AI Systems

arXiv: https://arxiv.org/abs/2505.10183 · HTML: https://mlsanity.com/paper/2505.10183/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 5 of its 15 references that are in the map.
- First posted on arXiv 15 May 2025.

## Builds on (most cited first, 5 of 5)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [MobileNetV2: Inverted Residuals and Linear Bottlenecks](https://mlsanity.com/paper/1801.04381/) (2018)
- [TensorFlow: Large-Scale Machine Learning on Heterogeneous Distributed Systems](https://mlsanity.com/paper/1603.04467/) (2016)
- [HetPipe: Enabling Large DNN Training on (Whimpy) Heterogeneous GPU Clusters through Integration of Pipelined Model Parallelism and Data Parallelism](https://mlsanity.com/paper/2005.14038/) (2020)
- [DaDu-E: Rethinking the Role of Large Language Model in Robotic Computing Pipeline](https://mlsanity.com/paper/2412.01663/) (2024)

## Built on by (most cited first, 1 of 1)

- [HetCCL: Accelerating LLM Training with Heterogeneous GPUs](https://mlsanity.com/paper/2601.22585/) (2026)

## Abstract (arXiv)

Embodied Artificial Intelligence (AI) systems, such as autonomous robots and intelligent vehicles, are increasingly reliant on diverse heterogeneous accelerators (e.g., GPGPUs, NPUs, FPGAs) to meet stringent real-time processing and energy-efficiency demands. However, the proliferation of vendor-specific proprietary communication libraries creates significant interoperability barriers, hindering seamless collaboration between different accelerator types and leading to suboptimal resource utilization and performance bottlenecks in distributed AI workloads. This paper introduces KAITIAN, a novel distributed communication framework designed to bridge this gap. KAITIAN provides a unified abstraction layer that intelligently integrates vendor-optimized communication libraries for intra-group efficiency with general-purpose communication protocols for inter-group interoperability. Crucially, it incorporates a load-adaptive scheduling mechanism that dynamically balances computational tasks across heterogeneous devices based on their real-time performance characteristics. Implemented as an extension to PyTorch and rigorously evaluated on a testbed featuring NVIDIA GPUs and Cambricon MLUs, KAITIAN demonstrates significant improvements in resource utilization and scalability for distributed training tasks. Experimental results show that KAITIAN can accelerate training time by up to 42% compared to baseline homogeneous systems, while incurring minimal communication overhead (2.8--4.3%) and maintaining model accuracy. KAITIAN paves the way for more flexible and powerful heterogeneous computing in complex embodied AI applications.