08/04 Seminar_ TPU V4

Jouppi, N., Kurian, G., Li, S., Ma, P., Nagarajan, R., Nai, L., ... & Patterson, D. A. (2023, June). Tpu v4: An optically reconfigurable supercomputer for machine learning with hardware support for embeddings. In Proceedings of the 50th Annual International Symposium on Computer Architecture (pp. 1-14). 의 내용이며, 모든 이미지의 출처는 해당 논문이다.

# Background

Market Share가 DLRM, Transformer(Large Language Models)의 비중이 큰 추세이다. 이에 맞춘 특성인 OCS, SparseCore를 사용하여 해당 특성을 가진다.

# Construction of V4

total 4096 TPU chip들을 사용한다.

16개의 CPU, 64개의 TPU가 하나의 rack이다. 아래는 total 8개의 rack array이다. 이런 rack array가 8개 더 있어야 4096개의 chip을 구성할 수 있다. rack 내부는 모두 electrical wire로 연결되어 있으며 서로 다른 rack들끼리는 ICI인 Optical link로 연결되게 된다.

# Optical Circuit Switching

1개의 rack은 64개의 TPU, 칩들은 4 x 4 x 4의 형태로 구성된다. 6방향 모두 link를 가져야하지만, optical wire가 양방향 통신을 지원하기 때문에, 3차원으로만 link를 가지면 된다. 해당 link 들을 switching하는 circuit을 통해서 어떤 block을 pass하여 연결할지를 정할 수 있고, 이를 통해서 TPU topology를 reconfiguration이 가능하다.

각 axis 별로 16개씩의 OCS가 필요하며, 양방향 통신이므로 3개의 면이 필요하다. 즉 total 48개의 OCS가 하나의 4 x 4 x 4 TPU block에 필요하다. 이 switch를 토대로 필요한 다른 block들 간의 연결을 할 수 있고 이를 통해서 failed unit을 pass하면서 computing performance를 높였다.

위 그림은 twisted torus 구조를 설명하는 예시이다. OCS를 통해서 해당 link들의 연결을 비교적 멀리 떨어져있더라도 direct하게 연결할 수 있다.

# Sparse Core

DLRMs(Deep Learning Recommendation Models)에서 사용되는 Embedding 작업에 특화된 H/W이다. Embedding은 추천 시스템에서의 input을 vector로 바꾸는 작업이다. 유튜브 영상을 embedding 작업을 통해서 어떤 category와 연관이 있는지의 정도를 숫자로 나타낼 수 있다.

보통은 DLRM의 첫번째 layer가 Embedding 작업을 수행한다.

SparseCore는 해당 Embedding 작업만을 위해서 만들어진 core라고 볼 수 있다. 해당 부분에 대한 정보가 자세히 안 나와있지만, Tensor core가 아닌 별도의 Core인 Sparse Core로 해당 작업을 수행하며, DMA, Sort, Sparse Reduce, Fork(Sequence와 관련된 분기문), Concat 기능들을 제공하며 Fetch, scVPU, Flush 등의 동작을 통해서 얼마만큼 해당 category와 관련이 있는지 vector로 나타낼 수 있다.

HBM, ICI를 활용하여 Supercomputer 전체 memory를 활용하여 해당 작업을 수행할 수 있다. facebook 같은 경우에는 해당 작업을 near memory로 동작하게끔 해서 더 속도를 높인다는데, 해당 부분이 CMEM으로 언급만 되어 있는 정도라 명확한 상관관계를 확인할 수 없었다.

#Performance comparison

TPU v4, nvidia A100, Graphcore IPU MK2의 성능 비교이다. 위 그래프는 BERT 모델을 training할 때, 아래 그래프는 ResNet을 training할 때이다. TPU와 달리 다른 두 custom ASIC 칩들은 멀티코어, 멀티 쓰레드 방식으로 NN연산을 지원한다. TPU v4의 경우 4096 chip을 전부 사용했을 때 BERT, ResNet 모두에서 가장 좋은 초당 연산율을 보여주었다.

Power 소모도 A100이 TPU v4에 비해서 1.33~1.93배 더 사용하고 있다.