TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hard

https://arxiv.org/abs/2304.01433
 
2020년에 구글 내부에서 출시된 TPU v4에 대한 논문이다. 계속 커져가는 LLM 모델 크기를 감당하기 위해 TPU v2가 256 node였던 것에 비해 TPU v4는 4096 node까지 확장할 수 있도록 설계되었다.
 
TPU는 사실상 GPU에 대항할 수 있는 거의 유일한 DSA라고 생각하고 있고, 게다가 이런 논문은 대규모 슈퍼컴퓨터를 실제로 운용해봐야 쓸 수 있는 주제라서 대기업 말고 쓸 수 있는 곳도 드물다. AI 하드웨어를 연구하고 싶다면 꼭 한번씩 읽어보면 좋겠다.
 
이 논문에서 소개하는 건 크게 TPU v4에서 새로 도입된

• Optical Circuit Switches (OCSes)
• 임베딩 지원을 위한 SparseCore 아키텍처
• All-to-All 통신을 위한 슈퍼컴퓨터 규모 토폴로지

이고 DSA 설계를 위해 생각해볼 점들이 후반부에 있다. 

---

2. Reconfiguarable Optical Switch

TPU v2에서 v3로 넘어갈 때 칩 개수가 4배 늘어났고, v3에서 v4로 넘어갈 때도 4배 늘리는 것을 목표로 했다.

2x4 2D Torus의 예시

기존 TPU들은 2D torus 토폴로지를 사용했는데, 칩 개수가 늘어날수록 한 쪽 끝에서 다른 쪽 끝을 연결하는 wrap-around channel이 지나치게 길어진다는 우려가 있었다.

4x4x4 Torus topology

때문에 v4에서는 3D torus 토폴로지로 전환했고 랙 내부는 electrical cable, global 연결은 optical fiber를 사용하는 방식을 채택했다. 

• 랙 내부에는 크기와 bisection bandwidth를 고려해서 4x4x4, 즉 64개의 TPU 칩을 전기 케이블로 연결한다. 
• 랙 간 연결은 optical link와 OCS를 사용한다.
• 정육면체 한 면에 16링크, 즉 블록 하나 당 96개의 optical link가 x축, y축, z축에 각각 16개씩 있는 OCS와 연결된다. 
• 64개 블록 x 64개 칩이므로 4096개의 TPU가 한 슈퍼컴퓨터를 구성하게 된다. 

 

OCS가 일반 전자식 스위치에 비해 가지는 장점은?

• 단일 광섬유로 양방향 통신 가능 -> 케이블 개수 절반
• Availibity: 장애가 발생한 링크나 호스트를 유동적으로 우회 가능
• Deployment: 일부 블록만 사용 가능하기 때문에 테스트가 용이
• Scheduling: v3에선 스케줄러가 연속된 256칩을 찾아야 했지만 v4에선 어떤 위치든 4개의 블록만 찾으면 OCS로 연결 가능 
• Modularity: 밀리초 단위로 토폴로지 변경 가능, OCS를 사용하면 토러스를 잘라도 메시가 아닌 토러스로 유지가 가능하다 -> Bisection 대역폭도 유지된다. 
• Security: 슬라이스 간에 네트워크 차원에서 격리가 가능하다. 
• Data, Tensor, Pipeline Parallelism을 사용할 때 물리적인 토폴로지를 최적화할 수 있다. 예를 들어 PP의 경우 4x4x32로 길쭉한 모양의 네트워크를 구성할 수 있고, TP나 DP를 주로 사용할 경우 8x8x8 정육면체 모양으로 구성하여 병렬화 전략에 맞게 토폴로지를 만들 수 있다.

4x2 TPU를 일반 torus와 twisted torus로 구성한 그림

또한 OCS의 유연성에 맞춰 논문에서는 Twisted Torus라는 토폴로지를 도입할 것을 제안했다. 일반 3D Torus는 worst-case latency가 큰데 twisted torus를 사용하면 throughput과 latency 모두 개선 가능하다. 물리적 배선을 직접 바꿀 필요 없이 OCS로 프로그래밍 가능하다. 
 
실제 사용 통계

• 전체 작업의 29%가 4x4x4 블록보다 작은 크기를 사용한다 -> 블록 내부는 전기식 스위치이기 때문에 2D mesh를 사용해야 한다. 
• 나머지 더 큰 작업들은 3D 토폴로지 사용이 가능하다. 
• n x n x 2n 또는 n x 2n x 2n (n이 4 이상일 때)인 경우에 twisted torus 구성이 가능하다. 
• 33% (71% 중 48%)가 twisted 가능한 조건을 만족했고 실제로 twisted tori를 사용한 경우는 28% (33%의 86%)였다. 
• 다르게 말하면 한 블록보다 크거나 같은 작업의 40%는 twisted tori를 사용한다. 
• OCS 비용은 전체 컴퓨터의 5% 미만이고 전력 소모도 전체의 3% 미만이라 도입 비용에 비해 이점이 크다. 

3. SparseCore: Embeddings Support

Deep Learning Recommendation Models (DLRM, 딥러닝 추천 모델)은 2022년 기준으로 전체 DL workload의 1/4 정도를 차지한다. 추천 모델은 광고, 검색 순위, 유튜브, 구글 플레이 같이 특정 사용자에게 맞는 정보를 제공하는 모델이다. DLRM 모델 크기는 예측 퀄리티, 학습 시간, 총 학습 비용, 서빙 지연, 그리고 총 서빙 비용에 의해 결정된다. 
 
DLRM의 입력은 거의 categorical feature로 이루어져있다. 각 categorical feature가 N 개의 이산값을 가지고 (N은 보통 vocab size) 일반적으로 굉장히 sparse하다. (전체 영어 단어 중에 사용자가 검색한 단어는 극히 일부일테니)
신경망이 효과적으로 학습하려면 단어를 dense vector로 바꿔야 한다. 이를 임베딩이라고 하며 이때 embedding lookup table을 사용하게 된다. 
 
단어 80000개를 임베딩 차원 100개로 바꾸는 table만 해도 8M개 원소가 필요하다. 크기 때문에 분산 학습이 필수적이다. table을 column 또는 row로 나눠 여러 칩에 저장하거나 테이블 별로 저장하는 방법이 있다. 이는 신경망으로 치면 model parallelism과 비슷하다. 작은 embedding table은 data parallelism도 성능 면에서 좋을 것이다. 
 
임베딩 연산의 특징은

• dense 연산(e.g. transformer, fully connected nn)과 다르게 작은 gather, scatter 메모리 접근이 많아 arithmetic intensity가 낮다. 그 말은 메모리 대역폭이 더 큰 병목이 된다는 뜻이다.
• 병렬화 방식에 따라 통신 패턴 역시 다르다. Model parallelism에서는 variable-length all-to-all exchange가 발생하고 bisection-bandwidth에 의해 속도가 제한된다.  DP에서는 AllReduce 통신이 발생하고 injection bandwidth에 의해 속도가 제한된다. 
• 입력이 sparse하기 때문에 특정 칩에 데이터가 몰리는 현상이 심하다. 자주 쓰이는 feature는 복제해서 메모리 접근과 통신량을 줄이는 것이 좋다. 

 
TPU의 TensorCore는 넓은 VPU와 matrix unit으로 이루어져있고 dense 연산에 특화되어있다. 임베딩 같이 sparse하고 메모리 대역폭 지향적인 연사을 TensorCore로 처리하면 비효율적이다. 또 CPU DRAM에 임베딩 테이블을 두게 되면 TPU-CPU 간 병목이 생겨 학습 속도가 급격하게 저하된다. 

때문에 임베딩 연산 및 HBM 활용을 위해 TPU v4에는 SparseCore (SC) 아키텍처가 추가되었다. SparseCore의 특징은 다음과 같다. 

• 칩 다이 면적과 전력의 5% 이하만 차지하지만 임베딩 성능은 5~7배 증가시킨다. 
• 슈퍼컴퓨터 전체 HBM을 하나의 메모리 공간으로 사용할 수 있다. 
• 16개의 코어(tile)들이 HBM 채널 각각에 직접 연결되어 고속 gather/scatter를 수행한다. 
• 타일마다 프로그래밍 가능한 8-wise SIMD VPU가 있다. 
• HBM에 업데이트된 파라미터를 쓰는 Flush Unit이 있다. 
• 5개의 Cross-Channel Unit이라는 모듈이 여러 tile에 걸친 공통 임베딩 연산을 수행한다. CISC 스타일의 명령어 집합을 지원하기 때문에 가변 길이 입력을 잘 처리할 수 있다. 

SC 유무에 따른 구글 내부 DLRM 워크로드의 성능

4. Using ML to Tailor the DNN to the TPU and the TPU Topology to the DNN

DNN의 정확도와 속도를 파레토 최적화 (한쪽이 이득을 보려면 다른 쪽이 반드시 손해를 보는 지점이라고 한다) 하기 위해서 Platform-Aware Neural Architecture Search (PA-NAS)를 개발했다. TPU v4 환경을 고려해서 모델 구조를 탐색하는 기법이다. 
PA-NAS를 사용해서 CNN1을 설계한 결과 baseline보다 1.6배 성능이 좋아졌고(QPS, latency는 개선하면서 정확도는 동일하게), DLRM을 설계한 결과 기존 모델은 SC가 25% 정도 시간에서 idle 상태였지만 개선 후에는 성능이 10% 정도 증가했다. 전문가가 수개월 동안 튜닝해야 얻을 수 있는 성능 향상이라고 논문에서는 평가했다. 

모델 구조 설계에 따라 throughput이 천차만별로 달라지는 모습

 
 
5~6은 벤치마크라 생략...
 

7. Discussion

토론할 점을 많이 제시하는 논문은 언제나 읽을 때 감사하다. 

7.1 Do peak FLOPS/second predict real performance?

TPU v3, 4, A100에서의 roofline model, 괄호 안의 값은 operational intensity

A100이 TPU v4보다 peak FLOPS/sec은 1.13배 더 높지만 실제 MLPerf 벤치마크에서는 TPU가 1.15~1.67배 더 빠르다. 언제나 암달의 법칙을 기억해야 한다. 루프라인은 최대 연산 성능과 메모리 대역폭만 고려한다. 칩 자체 성능뿐만 아니라 시스템 수준의 칩 간 인터커넥트, 그리고 전력이나 냉각 같은 요소까지 모두 고려해야 정확한 성능을 알 수 있다. 
 

7.2 How does OCS differ from NVLink and NVSwitch

• 모든 TPU는 자체적으로 빌트인 라우터를 가지고 있다.
• NVSwitch는 TPU의 라우터처럼 glueless한 연결을 가능하게 해주지만 스케일은 작다. (4~8 GPU vs 1024 TPU v3) 

(ICI: Inter-Core Interconnect)

• OCS는 MEMS 거울로 빛 신호를 반사만 해주기 때문에 입력과 출력이 1:1이다. WDM 방식으로 대역폭이 수 Tb/s까지도 상승할 수 있다. 
• NVSwitch는 전기적 패킷 스위칭이고 8GPU를 한 개의 스위치로, 256GPU를 2-level 스위치로 연결한다. 

7.3 What if TPU v4 used IB versus OCS?

• OCS 48대 (각각 128포트)를 40포트 IB 스위치로 대체하려면 568대가 필요하다.
• OCS와 포트당 가격은 비슷하지만 대역폭과 전력 면에서는 자체개발한 OCS가 더 낫다.
• IB와 ICI를 함께 쓸 때 소프트웨어 복잡도가 증가하는 것도 고려해야 한다.
• 또한 OCS 자체의 장점(2장에서 나왔던) 때문에라도 IB보단 OCS가 낫다고 평가했다. 

7.4 Nvidia announced the H100, the successor to A100, in 2022. Why not compare TPU v4 to it?

• 동시대 동공정 비교를 위해서 그러니까 빨리 다음 논문 내주시죠

7.5 Why 30~90% more power for A100?

완벽한 정량 분석은 어렵지만 TPU는 A100보다 on-chip SRAM이 4배(160MB vs 40MB) 크다. 덕분에 더 큰 블록 전송이 가능해서 에너지 효율이 좋다고 추측된다. 실제로 SRAM 사용을 32MB에서 160MB로 늘리면 성능/전력이 1.24배 향상되는 것을 관찰했다. 
또 세 가지 정성적인 요소들이 있다. 추가 연구가 필요하지만...

• 레지스터 파일 크기가 GPU가 100배 더 크다(0.25 MiB vs 27 MiB). 용량이 커지면 전력은 용량의 제곱근에 비례하므로 아마... 
• TPU v4의 128x128 MXU는 128 entry input을 128번 재사용한다. A100의 4x4 FP16 array는 4번 재사용된다. 여기서 on-chip SRAM 접근 횟수가 크게 달라질 수 있다. 
• A100이 다이 크기가 40% 더 크다. 통신에 필요한 에너지가 다를 것이다. 

7.6 What is the CO2e from TPU v4 vs other DSAs?

4M 계수(Model, Machine, Mechanization, Map)를 사용해서 계산해볼 수 있다. 
 
1. TPU v4와 다른 DSA에서 같은 모델을 학습시킨다고 가정하자. Model 계수를 1.0으로 놓는다. 
2. Performance/Watt로 측정한 값을 Machine 계수라고 하자. TPU v4는 A100보다 1.2~1.7배 빠르고 1.3~1.9배 전력 소모가 적다. 그래서 TPU의 Machine 계수는 2~6배 정도로 놓을 수 있는데 여기서는 보수적으로 2로 놓는다. 
3. 데이터센터의 효율을 PUE라고 하고 Mechanization 계수로 놓는다. 세계 평균 PUE는 1.57이고 구글 클라우드의 PUE는 1.10이다. 
4. 전력 그리드의 청정도를 CFE와 배출계수로 계산하고 이를 Map 계수로 놓는다. 미국 평균 CFE는 40%, 구글 오클라호마 데이터센터는 88%이다. 글로벌 평균 전력 배출 계수는 0.475 kg/kWh, 구글에서는 재생에너지를 사용하기 떄문에 0.074 kg/kWh다. 
 
다 계산해보면 CO2e는 다른 데이터센터보다 18.3배 적다. 와!
 

7.7 How fast do ML workloads change?

꽤 인상깊은 표라 다시 가져왔다.

• 2년 동안 RNN 비중이 29%에서 2%까지 내려가고, 그 사이에 Transformer가 57%까지 상승한 걸보자.
• BERT 논문 출판 2년 후에 BERT가 TPU v4에서 구글 워크로드의 25%를 초과했다.
• GPT-3 논문 출판 후에 LLM 워크로드 비중이 30%를 초과했다. 

슈퍼컴퓨터 설계, 구축, 배포에 2년 정도 시간이 걸리는 걸 생각하면 그 사이에 워크로드가 급격하게 바뀔 수 있다는 걸 하드웨어 설계자들이 반드시 알고 있어야 한다. 
 
하드웨어를 설계할 때는 유연성, 범용성을 반드시 확보해야 한다. 
 

7.8 Do individual DNN models also change?

DLRM0 모델의 크기 변화, 2020년에 무슨 일이 있었는지 궁금하다.

• 구글 내부의 DLRM0 모델은 2017년부터 2022년까지 가중치는 4.2배, 임베딩은 3.8배 커졌다.
• 5년 동안 새로운 버전이 약 6주마다 출시되었다.(총 43번)
• 모델 자체가 이 정도 속도로 바뀌니 하드웨어 기능, 즉 DSA 특성을 잘 활용하는 컴파일러가 필수적이다. 

 

7.9 Is MLPerf's DLRM benchmark realistic?

벤치마크 내용은 생략...

7.10 TPU v4 has less HBM capacity than A100; could that limit LLM performance?

• PA-NAS가 HBM 용량, ICI로 집계된 FLOPS/HBM 대역폭, 그리고 다른 모델 파라미터들(배치 사이즈 등)을 고려한다.
• HBM 용량이 제약으로 작용하는 경우도 있겠지만 일반적으로 칩 수를 더 늘리고, 계산-통신 overlap을 잘 활용하면 해결 가능하다. 
• A100이 HBM 용량이 더 크긴 하지만 NVLink로 직접 묶이는 도메인은 TPU보다 더 작다.
• 최적 설정은 GPU와 TPU 사이가 다를 수 밖에 없고 단순히 메모리 용량으로 결론 내리기 어렵다. 

7.11 How can DSAs avoid overspecialization?

ML은 워낙 급격하게 바뀌는 분야니 domain 특화, 범용성, 유연성 사이에 균형을 반드시 잘 잡아야한다. RNN 가속기를 2년 전에 만들기 시작했다고 생각해봐라... TPU는 유연한 설계 덕분에 배포를 시작한 후에 LLM이 급부상했는데도 그대로 사용할 수 있었다. 
 

8. Related Work

여러 가지 관련 연구들...

• TPU는 전통적인 슈퍼컴퓨터들처럼 3D torus지만 엔비디아 GPU는 2-tier fat-tree를 사용한다. 
• Twisted tori는 예전에도 ILLIAC-IV 등에서 사용되었다. 
• Shalf et al.과 Kamil et al.의 연구에서는 회선 교환과 패킷 교환을 혼합해서 사용했다. 
• TPU 말고도 OCS를 사용한 슈퍼컴퓨터가 많으니 참고하자. 상용 대규모 시스템에서 적용해서 재구성 가능한 인터커넥트로 성능을 향상한 사례는 TPU v4가 최초이다. 
• TPU v4는 물리 칩 전반에 논리적으로 공유된 메모리 주소를 가지고 있다. 소프트웨어가 접근/이동을 명시적으로 제어하고, 원격 메모리는 비동기 DMA '쓰기'만 허용한다. 과거에 Cray T3E는 유사하게 논리적으로 공유된 주소를 제공했는데, 추가로 벌크 비동기 읽기/쓰기, load-store, atomic operation까지 지원했다. 
• On-ASIC 임베딩 가속은 구글이 최초이다. Meta에서 Neo라는 GPU 시스템으로 임베딩 테이블을 학습한 사례가 있는데 테이블/행/열/데이터 병렬화를 모두 사용했고 통신-연산 overlap, 커널 퓨전까지 사용했다. 엔비디아 MLPerf에서도 비슷한 방식을 제시했다(인피니밴드에서 reduction을 위한 커스텀 커널 퓨전 등등)

---

시스템 전체의 균형을 맞춘다는 건 생각보다 훨씬 복잡하다. 메모리, 네트워크, 프로세서... 따로 파고들어 보면 명확해보이기도 하지만, 모든 요소가 한데 얽힐 때 어디서 병목이 터질지, 어떤 설계가 오랫동안 살아남을지 예측하기도 어렵고 모든 최선의 수를 써도 전혀 다른 결과가 나오기도 한다.
 
그럼에도 불구하고, 우리는 늘 답이 없는 문제를 해결하려고 한다. 이 과정 자체가 사실 연구의 본질이고, 그래서 재미있다.

복잡하고 답이 없는 시스템을 조금이라도 더 이해해보려 하고, 다음 번에는 더 나은 해법을 시도해보려 고민하는 이 과정. 그 자체가 논문을 읽고, 실험하고, 또 연구하게 만드는 이유가 된다.
 
그래서 결국 남는 건 한 문장이다. 
 
언제나 그렇지만 배울 점이 많은 논문을 써주셔서 감사합니다.