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

TL;DR

TEMP tackles the memory bottleneck of LLM training on wafer-scale chips by jointly co-designing topology-aware tensor-stream partitioning, traffic-aware mapping, and a wafer-specific search. The framework (TATP, TCME, DLWS) aligns compute and communication with the physical wafer topology to eliminate tensor replication, reduce tail latency, and mitigate inter-die contention. Through extensive evaluation, TEMP delivers substantial throughput improvements and enhanced power efficiency across multiple models, while providing robust fault tolerance and rapid cost-model-driven search. This work enables scalable, memory-efficient LLM training on wafer-scale fabrics, potentially surpassing traditional GPU clusters in performance and efficiency for large-scale deployments.

Abstract

Large language models (LLMs) demand significant memory and computation resources. Wafer-scale chips (WSCs) provide high computation power and die-to-die (D2D) bandwidth but face a unique trade-off between on-chip memory and compute resources due to limited wafer area. Therefore, tensor parallelism strategies for wafer should leverage communication advantages while maintaining memory efficiency to maximize WSC performance. However, existing approaches fail to address these challenges. To address these challenges, we propose the tensor stream partition paradigm (TSPP), which reveals an opportunity to leverage WSCs' abundant communication bandwidth to alleviate stringent on-chip memory constraints. However, the 2D mesh topology of WSCs lacks long-distance and flexible interconnects, leading to three challenges: 1) severe tail latency, 2) prohibitive D2D traffic contention, and 3) intractable search time for optimal design. We present TEMP, a framework for LLM training on WSCs that leverages topology-aware tensor-stream partition, traffic-conscious mapping, and dual-level wafer solving to overcome hardware constraints and parallelism challenges. These integrated approaches optimize memory efficiency and throughput, unlocking TSPP's full potential on WSCs. Evaluations show TEMP achieves 1.7x average throughput improvement over state-of-the-art LLM training systems across various models.

Figures (21)

• Figure 1: Illustrations for co-design features of the TEMP framework.
• Figure 2: Challenges and the TEMP solution for tensor-stream execution on wafer-scale chips. (a) Implementation challenges. (b) Key contributions and components of the TEMP framework.
• Figure 3: Hardware configurations of WSCs.
• Figure 4: (a) Illustration of 4-way tensor parallelism, where colors denote tensor partitions and numbers indicate device IDs. (b) Training time breakdown for GPT-3 and DeepSeek models with Megatron-LM. (c) Memory overhead of Megatron-LM compared to an ideal baseline. The dashed red line marks the WSC's per-die memory capacity.
• Figure 5: Challenges for deploying TSPP on WSC. $D_i$ means Die$i$ on a WSC.