Theseus: Exploring Efficient Wafer-Scale Chip Design for Large Language Models

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Zhu, Jingchen, Xue, Chenhao, Chen, Yiqi, Wang, Zhao, Zhang, Chen, Shen, Yu, Chen, Yifan, Cheng, Zekang, Jiang, Yu, Wang, Tianqi, Lin, Yibo, Hu, Wei, Cui, Bin, Wang, Runsheng, Liang, Yun, Sun, Guangyu
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866929600026640384
author Zhu, Jingchen
Xue, Chenhao
Chen, Yiqi
Wang, Zhao
Zhang, Chen
Shen, Yu
Chen, Yifan
Cheng, Zekang
Jiang, Yu
Wang, Tianqi
Lin, Yibo
Hu, Wei
Cui, Bin
Wang, Runsheng
Liang, Yun
Sun, Guangyu
author_facet Zhu, Jingchen
Xue, Chenhao
Chen, Yiqi
Wang, Zhao
Zhang, Chen
Shen, Yu
Chen, Yifan
Cheng, Zekang
Jiang, Yu
Wang, Tianqi
Lin, Yibo
Hu, Wei
Cui, Bin
Wang, Runsheng
Liang, Yun
Sun, Guangyu
contents The emergence of the large language model~(LLM) poses an exponential growth of demand for computation throughput, memory capacity, and communication bandwidth. Such a demand growth has significantly surpassed the improvement of corresponding chip designs. With the advancement of fabrication and integration technologies, designers have been developing Wafer-Scale Chips~(WSCs) to scale up and exploit the limits of computation density, memory capacity, and communication bandwidth at the level of a single chip. Existing solutions have demonstrated the significant advantages of WSCs over traditional designs, showing potential to effectively support LLM workloads. Despite the benefits, exploring the early-stage design space of WSCs for LLMs is a crucial yet challenging task due to the enormous and complicated design space, time-consuming evaluation methods, and inefficient exploration strategies. To address these challenges, we propose Theseus, an efficient WSC design space exploration framework for LLMs. We construct the design space of WSCs with various constraints considering the unique characteristics of WSCs. We propose efficient evaluation methodologies for large-scale NoC-based WSCs and introduce multi-fidelity Bayesian optimization to efficiently explore the design space. Evaluation results demonstrate the efficiency of Theseus that the searched Pareto optimal results outperform GPU cluster and existing WSC designs by up to 62.8\%/73.7\% in performance and 38.6\%/42.4\% in power consumption for LLM training, while improving up to 23.2$\times$ and 15.7$\times$ for the performance and power of inference tasks. Furthermore, we conduct case studies to address the design tradeoffs in WSCs and provide insights to facilitate WSC designs for LLMs.
format Preprint
id arxiv_https___arxiv_org_abs_2407_02079
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Theseus: Exploring Efficient Wafer-Scale Chip Design for Large Language Models
Zhu, Jingchen
Xue, Chenhao
Chen, Yiqi
Wang, Zhao
Zhang, Chen
Shen, Yu
Chen, Yifan
Cheng, Zekang
Jiang, Yu
Wang, Tianqi
Lin, Yibo
Hu, Wei
Cui, Bin
Wang, Runsheng
Liang, Yun
Sun, Guangyu
Hardware Architecture
The emergence of the large language model~(LLM) poses an exponential growth of demand for computation throughput, memory capacity, and communication bandwidth. Such a demand growth has significantly surpassed the improvement of corresponding chip designs. With the advancement of fabrication and integration technologies, designers have been developing Wafer-Scale Chips~(WSCs) to scale up and exploit the limits of computation density, memory capacity, and communication bandwidth at the level of a single chip. Existing solutions have demonstrated the significant advantages of WSCs over traditional designs, showing potential to effectively support LLM workloads. Despite the benefits, exploring the early-stage design space of WSCs for LLMs is a crucial yet challenging task due to the enormous and complicated design space, time-consuming evaluation methods, and inefficient exploration strategies. To address these challenges, we propose Theseus, an efficient WSC design space exploration framework for LLMs. We construct the design space of WSCs with various constraints considering the unique characteristics of WSCs. We propose efficient evaluation methodologies for large-scale NoC-based WSCs and introduce multi-fidelity Bayesian optimization to efficiently explore the design space. Evaluation results demonstrate the efficiency of Theseus that the searched Pareto optimal results outperform GPU cluster and existing WSC designs by up to 62.8\%/73.7\% in performance and 38.6\%/42.4\% in power consumption for LLM training, while improving up to 23.2$\times$ and 15.7$\times$ for the performance and power of inference tasks. Furthermore, we conduct case studies to address the design tradeoffs in WSCs and provide insights to facilitate WSC designs for LLMs.
title Theseus: Exploring Efficient Wafer-Scale Chip Design for Large Language Models
topic Hardware Architecture
url https://arxiv.org/abs/2407.02079