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 CompressionComputer Vision: SegmentationConstraint Satisfaction and Optimization: Distributed CSP/Optimization

Application Category

Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphsSemantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMsUser 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