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.

错误:搜索内容不能为空,请输入英文关键词
错误:关键词超出字数限制,请精简
高级检索

Segmentation-Aware Optimization of Collective for Waferscale Chips

  • Qize Yang,
  • Jiaxin Liu,
  • Taiquan Wei,
  • Yuxin Jin,
  • Shouyi Yin,
  • Yang Hu

摘要

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.