[논문 리뷰] TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings

Norman P. Jouppi, George Thomas Kurian외 12명|arXiv (Cornell University)|2023. 04. 04.

Photonic and Optical Devices인용 수 19

한 줄 요약

tldr: TPU v4는 광학적으로 재구성 가능한 인터커넥트를 광학 회로 스위치를 사용하여 도입하고, 임베딩 중심 워크로드를 위한 SparseCores를 추가하며, ML 성능과 효율성을 높이기 위해 4096칩으로 확장된다.

ABSTRACT

In response to innovations in machine learning (ML) models, production workloads changed radically and rapidly. TPU v4 is the fifth Google domain specific architecture (DSA) and its third supercomputer for such ML models. Optical circuit switches (OCSes) dynamically reconfigure its interconnect topology to improve scale, availability, utilization, modularity, deployment, security, power, and performance; users can pick a twisted 3D torus topology if desired. Much cheaper, lower power, and faster than Infiniband, OCSes and underlying optical components are <5% of system cost and <3% of system power. Each TPU v4 includes SparseCores, dataflow processors that accelerate models that rely on embeddings by 5x-7x yet use only 5% of die area and power. Deployed since 2020, TPU v4 outperforms TPU v3 by 2.1x and improves performance/Watt by 2.7x. The TPU v4 supercomputer is 4x larger at 4096 chips and thus ~10x faster overall, which along with OCS flexibility helps large language models. For similar sized systems, it is ~4.3x-4.5x faster than the Graphcore IPU Bow and is 1.2x-1.7x faster and uses 1.3x-1.9x less power than the Nvidia A100. TPU v4s inside the energy-optimized warehouse scale computers of Google Cloud use ~3x less energy and produce ~20x less CO2e than contemporary DSAs in a typical on-premise data center.

연구 동기 및 목표

• 빠르게 진화하는 ML 모델과 워크로드에 대응하여 확장 가능하고 유연한 ML 슈퍼컴퓨팅 하드웨어의 필요성을 동기 부여한다.
• 광학 회로 스위치와 SparseCores 임베딩 가속기를 포함한 TPU v4 아키텍처를 설명한다.
• 이전 TPU 및 관련 DSA 대비 성능, 전력 및 에너지 효율 개선을 정량화한다.
• 대형 언어 모델 및 데이터센터 에너지 발자국에 대한 배치 시사점을 강조한다.

제안 방법

• 광학 회로 스위치(OCSes)를 활용하여 규모, 가용성, 활용도, 모듈성, 배치, 보안, 전력 및 성능을 위해 인터커넥트 토폴로지를 동적으로 재구성한다.
• 임베딩 중심 모델에 최적화된 데이터플로우 프로세서인 SparseCores를 통합하여 다이 면적과 전력의 약 5% 수준을 사용하면서 5x-7x 가속을 달성한다.
• TPU v4를 4배 더 큰 4096칩 슈퍼컴퓨터로 배치하여 TPU v3에 비해 전체 성능이 최대 약 10배 빨라지도록 한다.
• 유사한 시스템 규모에 대해 Graphcore IPU Bow 및 Nvidia A100과의 비교 성능 및 효율성 지표를 제공한다.
• Google Cloud 창고-규모 배치에서의 에너지 및 CO2e 영향과 일반적인 온프렘 데이터센터를 비교하여 평가한다.

실험 결과

연구 질문

• RQ1 OCS를 통한 광학 재구성이 ML 워크로드에 대한 인터커넥트 성능, 활용도 및 신뢰성에 어떤 영향을 미치는가?
• RQ2 임베딩 중심 모델 가속에서 SparseCores의 효과는 무엇이며 결과적인 면적과 전력 소모는 어느 정도인가?
• RQ3 대규모 배치에서 TPU v4가 TPU v3, Graphcore IPU Bow, Nvidia A100에 비해 성능, 전력 및 CO2e 이점은 무엇인가?
• RQ4 TPU v4가 대형 언어 모델 및 창고 규모 데이터센터의 배치 고려사항에 어떤 영향을 미치는가?

주요 결과

• OCSes를 갖춘 TPU v4는 규모와 활용도를 개선하기 위한 동적 인터커넥트 토폴로지 선택(예: twisted 3D torus)을 달성한다.
• SparseCores는 다이 면적과 전력의 약 5%를 사용하여 임베딩 중심 모델에 대해 5x-7x의 속도 향상을 제공한다.
• TPU v4는 TPU v3보다 2.1배 빠르고 성능/와트는 2.7배 향상시킨다.
• 4096칩 TPU v4는 TPU v3보다 전반적으로 약 10배 빠르고 유사한 규모의 Graphcore IPU Bow보다 약 4.3배-4.5배 빠르다.
• TPU v4는 일반적인 온프레미스 DSA에 비해 Google Cloud 창고 규모 배치에서 에너지 소비가 더 적고 CO2e 배출도 더 적다(대략 에너지는 3배, CO2e는 20배 정도 적다).