Balanced segmentation of CNNs for multi-TPU inference

2024-10-22

Journal of Supercomputing

AI Summary

To address load imbalance, on-chip memory bottlenecks, and excessive inter-TPU communication overhead caused by uneven CNN model partitioning in multi-Edge TPU collaborative inference, this paper proposes a computation–communication co-modeling dynamic segmentation algorithm. For the first time, it achieves load-balanced partitioning at the CNN layer granularity, integrating graph partitioning, Edge TPU hardware-aware modeling, and topology-aware communication optimization, complemented by a lightweight runtime scheduler. Evaluated on ResNet-50 and ViT-L, the method achieves >92% TPU utilization, reduces end-to-end latency by 37%, and attains 2.1× higher throughput than baseline approaches. Key contributions are: (1) a fine-grained, dynamic segmentation framework supporting heterogeneous Edge TPU clusters; and (2) a joint optimization paradigm that simultaneously maximizes computational efficiency and communication locality.

Technology Category

• Machine Learning: Learning on the Edge & Model Compression
• Computer Vision: Segmentation
• Constraint Satisfaction and Optimization: Distributed CSP/Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation

Problem

• Optimize CNN segmentation for multi-TPU inference
• Balance workload across TPUs to reduce imbalance
• Address memory access bottleneck in Edge TPU systems

Innovation

• Balanced CNN segmentation for multi-TPU inference
• Compiler-based pipelined inference implementation
• Profiled-based workload balancing across TPUs