Segmentation-Aware Optimization of Collective for Waferscale Chips
摘要
Rapid scaling of large language models (LLMs) has led to an explosive increase in computation and communication demands, prompting interest in waferscale chips (WSCs). However, it suffers from inefficient collective communication due to the large diameter of the mesh network. Previous work improves bandwidth utilization via packet segmentation and pipeline scheduling, yet overlooks the trade-offs induced by segmentation configuration. In this work, we introduce a segmentation-aware cost model that extends the Alpha-Beta model to capture the connection between packet size, die-to-die (D2D) bandwidth, and pipeline depth. Based on this model, we propose an optimization framework that identifies optimal segmentation configuration and a pipeline scheduling strategy tailored to the dynamics of D2D bandwidth introduced by segmentation. Our approach reduces communication overhead by 73.43% and improves collective efficiency by 21.12% over SOTA algorithms.