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

TL;DR

TEMP is a framework that enhances large language model training on wafer-scale chips by optimizing tensor partitioning and mapping to overcome hardware constraints, significantly improving throughput.

Contribution

It introduces the tensor stream partition paradigm and a topology-aware, traffic-conscious mapping approach tailored for wafer-scale chips, addressing memory and communication challenges.

Findings

Achieves 1.7x average throughput improvement over state-of-the-art systems.

Effectively manages memory and communication bottlenecks in wafer-scale chip environments.

Demonstrates scalability across various large language models.

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.