Reproduce “TEMP: A Memory Efficient Physical-aware Tensor Partition-Mapping Framework on Wafer-scale Chips” on OpenResearch

About this paper

Huizheng Wang, Taiquan Wei, Zichuan Wang, Dingcheng Jiang, Qize Yang, Jiaxin Liu, Jingxiang Hou, Chao Li and 3 more

The TEMP framework optimizes large language model training on wafer-scale chips by introducing topology-aware tensor partitioning and a traffic-conscious mapping engine. It achieves up to 1.69x faster training throughput and up to 82% lower peak memory usage compared to baselines, effectively leveraging WSC architecture while mitigating tail latency and traffic contention.

Key insights

• Effectively mitigating tail latency on WSCs requires co-designing tensor partitioning and execution orchestration to ensure single-hop physical transfers and enable fine-grained communication-computation overlap.
• A unified representation of diverse parallel strategies coupled with iterative, traffic-aware communication path optimization is critical for reducing network contention on rigid 2D mesh interconnects.
• Combining an accurate, fast DNN-based cost model with a hierarchical search algorithm can dramatically reduce the time required to find optimal parallelization and mapping strategies for complex LLM deployments on WSCs.

Summary by alphaXiv.