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

2025-12-16

AI Summary

Training large language models (LLMs) on wafer-scale chips (WSCs) faces critical challenges: severe memory constraints, high tail latency induced by 2D mesh topology, intense on-wafer device-to-device (D2D) traffic contention, and exponential explosion of the tensor mapping search space. To address these, this paper proposes the Tensor Streaming Parallelism Paradigm (TSPP), the first topology- and traffic-aware end-to-end tensor parallel compilation and mapping framework tailored to WSC physical characteristics. TSPP introduces a joint topology–traffic optimization framework integrating topology-aware tensor sharding, analytical D2D traffic modeling, and a two-level solver—coarse-grained chip-level placement followed by fine-grained micro-op scheduling. Evaluated across multiple LLMs, TSPP achieves 1.7× average throughput improvement, significantly reduces tail latency and D2D congestion, and enables efficient training at kilo-GPU-equivalent scale.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Planning, Routing, and Scheduling: Planning with Language Models
• Natural Language Processing: (Large) Language Models

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Systems and Infrastructure for Web, Mobile and WoT: Data management and stream processing for Web, mobile and wireless applications
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

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.

Problem

• Optimizes memory efficiency and throughput for LLM training on wafer-scale chips
• Addresses communication bottlenecks and memory constraints in 2D mesh topologies
• Reduces search time for optimal tensor parallelism strategies on WSCs

Innovation

• Topology-aware tensor-stream partition for memory efficiency
• Traffic-conscious mapping to reduce D2D contention
• Dual-level wafer solving for optimal design search