KAITIAN统一异构加速器通信，提升机器人等系统训练效率
KAITIAN: A Unified Communication Framework for Enabling Efficient Collaboration Across Heterogeneous Accelerators in Embodied AI Systems
- 用统一抽象层融合厂商优化与通用协议，实现跨设备协同
- 动态调度使训练提速最高达42%，通信开销仅2.8%~4.3%
- 适合需高效融合GPU、NPU等异构硬件的智能体AI研发团队
具身人工智能系统（如自主机器人、智能汽车）日益依赖多种异构加速器（如GPGPU、NPU、FPGA）以满足实时处理与能效要求。然而，各厂商专有通信库导致严重互操作性障碍，阻碍不同加速器间的协作，造成资源利用率低与性能瓶颈。本文提出KAITIAN，一种新型分布式通信框架，通过统一抽象层，智能集成厂商优化通信库以提升组内效率，并采用通用协议保障组间互操作性。其关键创新在于负载自适应调度机制，可依据设备实时性能动态分配任务。在集成NVIDIA GPU与寒武纪MLU的测试平台上，基于PyTorch实现的KAITIAN在分布式训练中显著提升资源利用率与可扩展性。实验表明，相较基线同构系统，训练时间最多缩短42%，通信开销仅为2.8%–4.3%，且模型精度保持不变。KAITIAN为复杂具身智能应用中的异构计算提供了更灵活高效的解决方案。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。