기타 대시보드

1. 왜 지금 스케일링을 알아야 하나

3~4년 전만 해도 대다수 머신러닝 연구자는 이 책에 나오는 내용을 몰라도 됐다고 저자들은 말한다. 지금은 "작은" 모델조차 하드웨어 한계에 바짝 붙어 돌아가서, 새로운 연구를 하려면 규모에서의 효율을 같이 생각해야 한다. 저자들이 짚는 과거 사례는 알렉스 크리제프스키다. CNN(합성곱 신경망)을 빠르게 돌리려고 날것의 CUDA 코드를 직접 짰던 그의 작업은 몇 년 뒤 Theano·TensorFlow 같은 라이브러리가 대신 처리해 줬다. 저자들은 지금 책에 담은 내용도 몇 년 안에 그렇게 추상화될 수 있다고 인정하면서도, 스케일링 법칙이 모델을 계속 하드웨어의 한계선까지 밀어붙이는 한 최전선 연구는 대규모로 모델을 효율적으로 돌리는 법과 떼어 놓을 수 없을 거라고 본다. 저자들의 표현을 빌리면 "벤치마크에서 20% 이기더라도 루프라인 효율에서 20%를 깎아 먹으면 의미가 없다." 유망한 모델 구조가 실패하는 이유도 대개 둘 중 하나다. ① 규모에서 효율적으로 못 돌거나, ② 그렇게 돌아가게 만드는 작업에 아무도 공을 들이지 않아서다.

2. 강한 스케일링이란 무엇인가

저자들이 "모델 스케일링"이라 부르는 목표는 단순하다. 학습이나 추론에 쓰는 칩 수를 늘릴 때 처리량도 그만큼 비례해서 늘리는 것, 이것이 "강한 스케일링(strong scaling)"이다. 칩을 더 붙이는 병렬화는 계산 시간을 줄여 주지만 그 대가로 칩 사이에 오가는 통신이 늘어난다. 통신에 걸리는 시간이 계산 시간을 넘어서면 "통신에 발목 잡힌(communication bound)" 상태가 되고, 그 순간부터는 칩을 더 붙여도 처리량이 비례해서 늘지 않는다. 계산 시간이 줄면 이번엔 칩 한 개 수준의 병목이 드러난다. 저자들은 500조 연산/초를 낸다고 표시된 TPU나 GPU라도 파라미터를 메모리에서 옮기는 데 발목 잡히면 표시치의 10분의 1만 낼 수 있다고 짚는다. 칩 하나가 처리하는 연산량과 메모리 대역폭, 전체 메모리 용량이 스케일링 이야기의 핵심에 있는 이유다. 이 병목이 어디서 나타날지 미리 알면 그것을 피하도록 모델을 설계하거나 다시 짤 수 있다는 것이 저자들의 논리다.

3. 하드웨어와 알고리즘이 함께 움직인 이야기

하드웨어를 설계하는 쪽은 반대편 문제를 짊어진다. 비용을 최소로 두면서 알고리즘이 딱 필요한 만큼의 연산·대역폭·메모리를 내주는 하드웨어를 만들어야 한다. 저자들은 이 코디자인(co-design, 하드웨어와 알고리즘을 함께 걸고 설계하는 문제)이 얼마나 위태로운지 짚는다. 실제 칩이 나오기까지 2~3년이 걸리는데, 그사이 알고리즘이 어떤 모습일지 미리 걸어야 한다. TPU는 이 내기에서 이긴 사례로 저자들이 꼽는 이야기다. 행렬곱은 메모리 바이트당 소화하는 FLOPs(부동소수점 연산 수)가 다른 어떤 연산보다도 많은(바이트당 N FLOPs) 독특한 알고리즘이고, 시스톨릭 배열(데이터가 격자 모양으로 늘어선 연산 유닛 사이를 리듬감 있게 흘러가며 계산되는 구조)로 지은 초기 TPU는 나온 시점의 GPU보다 달러당 성능에서 훨씬 앞섰다고 저자들은 말한다. TPU는 처음부터 머신러닝 작업에 맞춰 설계됐고, 텐서 코어를 얹은 GPU도 빠르게 같은 틈을 메워 가는 중이다. 반대로 신경망이 그때 뜨지 않았거나 TPU가 다루기 힘든 방향으로 근본적으로 바뀌었다면, GPU보다 유연성이 떨어지는 TPU에 건 그 내기는 값비싼 실패로 남았을 거라고 저자들은 짚는다.

4. 12장을 4부로 어떻게 엮었나

책은 4부 12장으로 짜여 있다. 1부(1~3장)는 예비지식이다. 1장은 계산·통신·메모리 세 가지가 알고리즘 속도를 어떻게 가두는지 다루는 루프라인(roofline, 무엇이 속도를 묶는지 재는 분석) 분석이고, 2장은 TPU가 칩 하나로서 그리고 제한된 대역폭·지연시간을 가진 칩 사이 연결로 묶인 시스템으로서 어떻게 동작하는지, 3장은 여러 TPU에 흩어진 행렬을 어떻게 곱하는지를 샤딩(sharding, 행렬을 여러 칩에 쪼개 나눠 담는 것)으로 설명한다. 이 세 장에서 저자들은 행렬곱 하나가 계산에 발목 잡히는지 메모리·통신에 발목 잡히는지, TPU가 어떻게 학습 클러스터로 배선되고 부분마다 대역폭을 얼마나 갖는지, 여러 TPU에 흩어진 배열을 모으고 흩뿌리고 다시 나누는 데 걸리는 시간과 서로 다르게 흩어진 행렬을 효율적으로 곱하는 법을 답한다. 2부(4~8장)는 트랜스포머다. 4장은 순전파·역전파에 드는 FLOPs, 파라미터 수, KV 캐시(어텐션에 쓰는 키·값을 저장해 두는 캐시) 크기까지 트랜스포머 수학을 다루고, 이 계산으로 모델이 메모리를 얼마나 쓰는지, 계산과 통신에 시간을 얼마나 쓰는지, 어텐션이 피드포워드 블록에 비해 언제 중요해지는지를 알 수 있다. 5장과 7장, 학습과 추론 장이 이 책의 중심이라고 저자들은 밝힌다. 모델 크기와 칩 수가 주어졌을 때 어떻게 나눠야 강한 스케일링 영역에 머무는지를 묻는 물음이다. 6장과 8장은 이 개념들을 LLaMA 3라는 널리 쓰는 오픈소스 모델에 적용하는 실습이다. 3부(9~10장)는 9장이 JAX+XLA 스택과 JAX/텐서보드 프로파일러로 실제 문제를 디버깅하는 법을, 10장이 계산을 병렬화하는 JAX API를 예제로 다룬다. 4부(11~12장)는 마무리 장과, GPU가 어떻게 동작하고 어떻게 배선되며 루프라인이 TPU와 어떻게 다른지 새로 다루는 12장이다.

5. 무엇을 쪼개고 무엇을 줄이나

모델을 여러 칩에 나눠 쓰는 데는 두 갈래 기법이 있다고 저자들은 짚는다. ① 계산을 쪼개는 병렬화 넷(데이터·텐서·파이프라인·엑스퍼트), ② 메모리 요구량 자체를 줄이는 기법 몇 가지다. 재계산(rematerialization, 중간 계산값을 저장하는 대신 필요할 때 다시 계산하는 것), 옵티마이저·모델 샤딩(ZeRO로 불리는 방식), 호스트 오프로드(파라미터를 호스트 메모리로 내보내는 것), 그레이디언트 누적이 후자에 들어간다. 5장과 7장에서 이 목록을 자세히 다루며, 주어진 칩 수와 모델 크기에서 어떤 조합을 골라야 강한 스케일링 영역에 머무는지를 저자들은 풀어낸다. 책을 처음부터 끝까지 순서대로 읽을 필요는 없다고 저자들은 밝힌다. 1~3장은 전제 지식과 뒤에서 쓸 표기를 세우는 자리라 이미 익숙하면 건너뛰어도 된다. 저자들은 책 끝에 제임스 브래드버리와 블레이크 헥트먼이 이 책에 담긴 여러 아이디어를 이끌어 냈다고 밝혀 둔다.

1. 계산은 얼마나 걸리나

저자들은 이 장을 단순한 물음 하나로 연다. 어떤 알고리즘이 왜 50초나 5밀리초가 아니라 50밀리초가 걸리는지, 모델 안에서 실제로 무엇이 그 시간을 채우는지를 묻는다. 답은 계산과 통신 둘로 갈리고, 이 갈래는 뒤에 이어지는 장에서도 그대로 되풀이해 쓰인다.

딥러닝 모델은 결국 행렬곱의 더미이고, 곱셈과 덧셈을 합쳐 FLOP(부동소수점 연산)이라 부른다. 계산 시간은 처리할 FLOPs를 가속기가 초당 처리하는 FLOPs로 나눈 값이다(`T_math = 계산 FLOPs / 가속기 FLOPs/s`).

NVIDIA H100은 bf16(16비트 부동소수점 형식) 연산을 초당 9.89e14회 처리하고, TPU v6e는 초당 9.1e14회 처리한다. 1e12 FLOPs짜리 계산이라면 H100에서 약 1.01밀리초, TPU v6e에서 약 1.1밀리초가 걸린다는 계산이 나온다. 다만 H100과 B200 같은 GPU는 공표된 최대치의 80~85%밖에 못 내는 경우가 많고, TPU는 실사용에서 95%에 가깝게 낸다고 저자들은 덧붙인다. 이 비교가 가격까지 맞춘 것은 아니라는 단서도 붙는다. H100과 TPU v6e는 가격이 다르므로 처리 속도만으로 어느 쪽이 유리한지 가릴 수는 없다.

2. 데이터를 나르는 시간도 함께 잰다

칩 안에서는 텐서를 고대역폭 메모리(HBM — 칩 옆에 쌓은 주 메모리)와 연산 코어 사이로 옮겨야 한다. H100은 이 대역폭이 초당 3.35테라바이트, TPU v6e는 초당 1.6테라바이트다. 모델을 여러 가속기에 나눠 실으면 텐서는 칩과 칩 사이도 오간다. 이때 쓸 수 있는 경로는 보통 ICI(칩 사이를 잇는 전용 연결)·DCN(데이터센터 네트워크)·PCIe 셋이고, 셋 다 대역폭이 다르다.

칩 안이든 칩 사이든 통신 시간은 같은 식으로 구한다. 옮길 바이트 수를 대역폭(초당 바이트)으로 나눈다(`T_comms = 통신 바이트 / 대역폭 바이트/s`). 계산과 통신은 대개 겹쳐 돌릴 수 있으므로 실행 시간의 하한은 둘 중 큰 값이고, 전혀 안 겹친다고 보수적으로 잡으면 상한은 둘을 더한 값이다. 저자들은 실무에서 이 하한을 기준으로 최적화한다고 밝힌다. 겹쳐 쓰기를 최대한 살리면 실제 시간이 하한에 가까워지기 때문이다. 최댓값을 기준으로 잡으면 하한과 상한의 차이도 묶인다. 상한(둘의 합)은 하한(둘 중 큰 값)의 최대 두 배를 넘지 않는다. 이보다 더 정확한 값이 필요하면 계산과 통신이 겹치는 구간과 그 밖의 오버헤드까지 따져야 하는데, 이는 실제 모델을 프로파일링해야 나오는 값이라고 저자들은 덧붙인다.

계산이 통신보다 오래 걸리면 연산기를 놀리지 않고 다 쓴다는 뜻이라 "연산 병목(compute-bound)"이라 부른다. 반대로 통신이 더 오래 걸리면 연산기 일부가 데이터를 기다리며 논다는 뜻이라 "통신 병목(comms-bound)"이라 부른다. 이 둘을 나누는 잣대가 산술강도다.

3. 산술강도라는 잣대

산술강도(arithmetic intensity)는 총 FLOPs를 통신 바이트 수로 나눈 값이다. 바이트 하나를 옮길 때마다 몇 번의 연산을 하는지를 재는 셈이다. 어떤 연산의 산술강도가 가속기 자체의 산술강도(최대 FLOPs/s를 대역폭 바이트/s로 나눈 값)보다 크면 연산 병목에 들어서고, 작으면 통신 병목에 걸린다.

TPU v5e의 행렬곱 유닛(MXU)은 초당 1.97e14 FLOPs를 처리하고 HBM에서 초당 8.2e11바이트를 읽어 온다. 둘을 나누면 240 FLOPs/byte가 나오고, 이것이 TPU v5e의 임계 산술강도다. 어떤 연산의 산술강도가 240보다 낮으면 바이트를 읽어 오는 속도가 시간을 잡아먹는다는 뜻이다.

이 비교는 세 줄로 정리된다. 계산 시간이 통신 시간보다 크다는 것은 계산 FLOPs를 가속기 FLOPs/s로 나눈 값이 통신 바이트를 대역폭으로 나눈 값보다 크다는 뜻이고, 양변을 정리하면 계산 FLOPs를 통신 바이트로 나눈 값(연산의 산술강도)이 가속기 FLOPs/s를 대역폭으로 나눈 값(가속기의 산술강도)보다 크다는 뜻으로 바뀐다. 두 산술강도만 견주면 실제로 돌려 보지 않고도 어느 쪽이 병목인지 미리 알 수 있는 이유가 여기에 있다.

이 잣대가 얼마나 낮게 나올 수 있는지 보여주는 예가 두 벡터의 내적이다. 길이 N짜리 bf16 벡터 두 개를 곱해 더하려면 2N바이트씩 두 번 읽고 N번 곱하고 N-1번 더한 뒤 2바이트를 다시 써야 한다. N이 커질수록 이 비율은 1/2로 수렴한다. 240은커녕 1에도 못 미치는 값이라 내적은 어떤 하드웨어에서도 거의 항상 통신 병목에 걸린다고 저자들은 말한다. 다만 저자들은 각주에서 내적이 실제로는 행렬곱 유닛(MXU)이 아니라 벡터 처리 유닛(VPU)에서 돈다는 점을 짚는다. TPU v5p의 VPU는 코어당 초당 약 7e12 FLOPs를 처리해 임계 산술강도가 약 3에 불과하다. 산술강도 1/2인 내적은 이 낮은 기준으로 봐도 여전히 통신 병목에 걸린다.

4. 루프라인 그래프로 읽기

이 관계를 그래프 하나로 정리한 것이 루프라인 도표다. 가로축에 산술강도, 세로축에 실제로 달성하는 FLOPs/s를 놓는다. 산술강도가 낮은 구간에서는 그래프가 직선으로 올라간다. 대역폭이 한계라 산술강도가 늘어난 만큼 처리량도 그대로 늘어나기 때문이다. 임계 산술강도를 넘어서면 그래프는 수평으로 꺾인다. 이미 연산기를 다 쓰고 있어서 산술강도를 더 올려도 처리량이 늘지 않는다. 저자들이 든 예시 그래프에서는 서로 다른 산술강도를 가진 두 알고리즘(Algo 1·Algo 2)을 두 가지 대역폭(BW1·BW2) 아래 나란히 놓는다. 대역폭을 BW1에서 BW2로 키우면 꺾이는 지점이 오른쪽 위로 옮겨 가면서 같은 산술강도에서도 더 높은 처리량을 낼 수 있다. 이 그래프는 구간을 셋으로 나눈다. 두 대역폭 모두에서 통신에 발목 잡히는 구간, 낮은 대역폭 BW1에서만 병목이 걸리는 구간, 어느 대역폭에서도 연산기를 다 쓰는 구간이다. 가운데 구간에 있는 알고리즘은 대역폭만 BW1에서 BW2로 올려도 병목에서 벗어난다.

5. 행렬곱은 배치 크기 하나로 정리된다

행렬곱 `X[B,D] × Y[D,F] → Z[B,F]`의 산술강도는 `2BDF / (2BD + 2DF + 2BF)`다. 트랜스포머의 행렬곱처럼 배치 크기 B가 D, F보다 훨씬 작다고 두면 이 식은 대략 B로 줄어든다. 그러면 산술강도가 240을 넘는다는 조건이 그대로 배치 크기가 240을 넘는다는 조건이 된다. TPU v5e에서 토큰 배치 크기가 240을 넘으면 그 행렬곱은 연산 병목에 들어선다는 단순한 규칙이다. GPU에서는 이 문턱이 300에 가깝다고 저자들은 덧붙인다.

저자들은 이 가정이 트랜스포머에 맞는 이유도 짚는다. 한 칩(레플리카)이 맡는 로컬 토큰 배치 크기는 보통 1,024보다 작은 반면 D와 F는 보통 8,000을 넘는다. 예를 들어 4,096토큰짜리 시퀀스 512개를 GPU 2,048개에 나눠 돌리면 전체 배치는 512×4,096으로 200만 토큰이지만, 칩 하나가 맡는 로컬 배치는 1,024토큰에 그친다. 여기서 배치는 시퀀스 개수가 아니라 토큰 개수로 잰다는 점도 저자들은 강조한다. 시퀀스가 같든 다르든 토큰 수만 맞으면 루프라인 계산은 동일하다. 이 임계 배치는 모델을 여러 칩에 나눠 싣더라도(샤딩) 그대로 적용되는데, 샤딩으로 칩 수를 늘리면 연산량과 대역폭이 함께 늘어나므로 임계 배치는 가중치 한 벌을 기준으로 정해지기 때문이라고 저자들은 짚는다.

6. 정밀도를 낮추면 문턱이 어떻게 움직이나

실제로 큰 행렬곱을 실행할 때는 행렬 전체를 한 번에 못 올리고 VMEM·SMEM 같은 칩 안 저장소에 맞는 작은 타일(bm·bk·bn 크기)로 쪼개 나눠 싣는데, 그러면 같은 데이터를 여러 번 읽게 되어 실제 통신량이 앞의 근사보다 늘어날 수 있다. 이때의 산술강도는 대략 bm·bn/(bm+bn)으로, 앞서 구한 근사와 같은 꼴로 정리된다고 저자들은 각주로 밝힌다.

정밀도를 int8로 낮추면 문턱이 어떻게 움직이는지도 저자들은 문제로 짚는다. 가중치와 연산을 모두 int8(1바이트 정수형)로 낮추면 산술강도가 2B로 커지지만, 하드웨어의 int8 처리량(TPU 기준 초당 3.94e14 OPs로 bf16의 약 2배)도 함께 커져 임계 배치는 그대로 240 근처에 남는다. 반면 가중치만 int8로 낮추고 활성화와 연산은 bf16을 유지하면 임계 배치는 120으로 절반이 된다. F=D=4096과 F=D=1024를 놓고 배치 크기별 처리량을 그려 보면 둘 다 결국 같은 하드웨어 최대치에 닿지만, D와 F가 작을수록 임계 배치는 커진다. D=F=1024는 D=F=4096보다 임계 배치가 거의 두 배다.

배치마다 다른 행렬을 곱하는 경우는 다르다. `int8[B,D] · int8[B,D,F] → int8[B,F]`처럼 B개의 서로 다른 [D]×[D,F] 곱을 한 번에 하면 총 연산량은 그대로 2BDF지만 통신 바이트는 BD+BDF+BF로 늘어나 BDF 항이 지배한다. 그러면 산술강도는 배치 크기와 상관없이 대략 2로 고정되어, 어떤 배치를 골라도 거의 항상 통신 병목에서 벗어나지 못한다.

GPU에서도 같은 계산이 성립한다. NVIDIA 스펙시트가 밝힌 H100의 bf16 처리량은 구조적 희소성을 전제로 한 1.979e15 FLOPs/s이고, 희소성을 안 쓰면 그 절반인 9.89e14 FLOPs/s다. 여기에 대역폭 3.35e12바이트/s를 나누면 임계 배치는 약 295로 나와, TPU의 240과 크게 다르지 않다고 저자들은 밝힌다.

7. 칩을 넘어가는 통신도 같은 잣대로 잰다

가속기 두 개에 행렬을 D 차원 기준으로 절반씩 나눠 곱할 때도 같은 계산이 적용된다. 각 칩은 전체 계산량의 절반만 맡으므로 계산 시간은 그대로 절반이 된다. 두 칩은 각자 구한 부분합을 상대 칩으로 보내 더해야 하므로, 통신 바이트는 2BF이고 통신 시간은 이를 칩 간 대역폭(저자들이 예로 든 초당 4.5e10바이트)으로 나눈 값이다.

이 경우 연산 병목이 되는 조건은 `D/2 > 4377`, 즉 `D > 8755`다. 앞서 배치 크기 B에 걸려 있던 조건이 이번에는 D에 걸린다는 점이 다르다. 저자들은 이 예시가 다소 인위적으로 고른 것이라고 인정하면서도, 이 책이 다루는 루프라인 대부분은 칩 안이 아니라 칩 사이 통신 쪽이라고 못박는다. 여러 TPU에 걸쳐 나뉜 행렬곱을 언제 병렬로 돌려도 되는지 판단하는 데 바로 이런 계산이 결정적이라는 것이다. 칩 안 통신이든 칩 사이 통신이든 같은 방식으로 하한과 상한을 구할 수 있다는 것이 이 장의 요지라고 저자들은 정리한다.

1. TPU는 행렬곱 기계다

TPU(Tensor Processing Unit)는 행렬곱에 특화된 연산 코어인 TensorCore에 고대역폭 메모리(HBM, High Bandwidth Memory — 칩 옆에 쌓아 붙인 큰 저장소)를 붙인 칩이다. 저자들은 TensorCore 안에서 셋을 짚는다. ① MXU(Matrix Multiply Unit)는 시스톨릭 배열(systolic array — 데이터를 격자 모양 회로 사이로 흘려보내며 계산하는 방식)을 써서 8사이클마다 `bf16[8,128]` 행렬과 `bf16[128,128]` 행렬을 곱해 `f32[8,128]` 결과를 낸다. TPU v5e는 1.5GHz에서 MXU 하나가 초당 `5e13`번의 bf16 연산을 하고, 칩 한 장에 MXU가 2~4개 있어 칩 전체로는 초당 `2e14`번(int8이면 `4e14`번)이다. 트릴리엄(TPU v6e)의 MXU는 256×256으로, v5e까지 써 온 128×128보다 커서 사이클당 4배 많은 FLOPs를 낸다. ② VPU(Vector Processing Unit)는 ReLU 같은 일반 연산과 벡터 사이의 덧셈·곱셈, 합산을 맡는다. ③ VMEM(Vector Memory)은 TensorCore 바로 옆에 둔 온칩 저장소로 v5e 기준 128MiB에 불과하지만 MXU로 가는 대역폭은 HBM보다 훨씬 높다. CPU의 L1·L2 캐시와 동작이 비슷하지만 훨씬 크고, 무엇을 올려 둘지는 프로그래머가 직접 정한다. HBM에 있는 데이터를 TensorCore가 쓰려면 먼저 VMEM으로 옮겨야 한다.

TPU v5p는 저자들이 "지금까지 나온 것 중 가장 강력한 축에 든다"고 부르는 세대다. 코어 하나가 초당 `2.5e14`번, 칩 하나가 초당 `5e14`번의 bf16 연산을 하고, 칩 8,960개를 묶은 포드 하나는 초당 4엑사플롭스(bf16)를 낸다. 이만한 컴퓨터를 구글은 여러 대 갖고 있다.

2. 파이프라인 — MXU를 놀리지 않는다

TPU의 모든 연산은 파이프라인으로 겹쳐 돈다. 행렬 X와 A를 곱해 Y를 낼 때 TPU는 먼저 A와 X의 조각을 HBM에서 VMEM으로 복사하고, 이를 MXU에 실어 X는 8×128 조각으로 A는 128×128 조각으로 곱한 뒤 결과를 다시 조각씩 HBM으로 돌려보낸다. 이 복사를 MXU 연산과 겹쳐 두면 MXU는 메모리 전송을 기다리지 않고 계속 돌아가고, 계산은 메모리가 아니라 연산량에 묶인다(compute-bound).

연산 강도(arithmetic intensity — 옮긴 1바이트당 몇 번 연산하는지)가 낮은 계산일수록 이 겹침의 값어치가 커진다. VMEM 대역폭은 HBM보다 약 22배 높아서, VMEM에서 읽고 쓰는 MXU 연산은 연산 강도가 10~20만 되어도 최대 성능에 닿는다. 반대로 가중치를 HBM에서 곧장 읽으면 같은 성능을 내는 데 훨씬 큰 배치 크기가 있어야 한다. 가중치를 VMEM에 미리 올려 두면(프리페치) 다음 계산의 적재 비용을 앞 계산 시간 뒤에 숨길 수 있지만, VMEM이 128MiB로 워낙 작아서 층 하나가 통째로 들어갈 만큼 나눠 두어야만 가능하다고 저자들은 짚는다.

3. 칩 한 장 안 — 메가코어와 트레이

TPU 칩 한 장은 보통(전부는 아니지만) 코어 두 개가 메모리를 나눠 쓰며 사실상 연산량이 두 배인 가속기 하나처럼 움직인다 — 저자들은 이를 메가코어(megacore) 구성이라 부른다. v4·v5·v6 세대가 이 구조를 쓰고, v7은 메가코어 대신 코어 두 개 사이에 고대역폭 링크를 따로 둔다. 더 오래된 v3 이전 칩은 코어마다 메모리가 따로 있어 별개의 가속기 두 개로 취급된다. 추론용으로 만든 v5e는 칩당 코어가 하나뿐이다.

칩은 넷씩 묶여 트레이(tray) 하나를 이루고, 트레이는 PCIe로 CPU 호스트에 붙는다. 흔히 보는 구성은 트레이 하나(칩 4장, 코어 8개, 보통 메가코어 4개로 다룸)를 코랩이나 TPU-VM 하나로 노출하는 형태다. v5e 같은 추론용 칩은 호스트 하나에 트레이가 둘이라 칩 8장이 곧 코어 8개다. PCIe 대역폭은 넉넉하지 않다 — v4 기준 방향마다 초당 16GB로 HBM보다 100배쯤 느리다. 호스트(CPU) 램으로 데이터를 옮기거나 되불러올 수는 있지만 빠르지는 않다.

4. 칩과 칩 사이 — ICI가 만드는 토러스

칩들은 ICI(Inter-Chip Interconnect) 네트워크로 포드(pod) 안에서 서로 이어진다. 오래된 세대(v2·v3)와 추론용 칩(v5e), 트릴리엄(v6e)은 ICI가 가장 가까운 4개 칩과 이어져 2차원 토러스(고리 모양으로 감아 끝과 끝을 이은 격자)를 이루고, v4와 v5p는 가장 가까운 6개와 이어져 3차원 토러스를 이룬다. 이 연결은 호스트를 거치지 않고 칩과 칩이 직접 잇는다. 고리 구조는 두 지점 사이 최대 거리를 N에서 N/2로 줄여 통신을 빠르게 하고, TPU는 여기에 토러스를 뫼비우스 띠처럼 한 번 더 비트는 "뒤틀린 토러스" 배선을 더해 평균 거리를 한층 줄인다.

ICI로 묶은 포드는 아주 커질 수 있다. 최대 포드 크기(슈퍼포드)는 v4가 16×16×16, v5p가 16×20×28이다. 이 큰 포드는 광학 배선으로 다시 연결할 수 있는 4×4×4 칩짜리 정육면체를 이어 붙여 만든다. v5e와 트릴리엄 포드는 16×16 2차원 토러스 하나로 끝이고 그 이상은 못 늘어나지만, 포드끼리는 데이터센터 네트워크(DCN)로 통신할 수 있다.

TPU와 GPU가 갈리는 지점이 이 연결 방식이다. GPU는 스위치 계층을 여러 겹 쌓아 사실상 모든 GPU 사이를 점대점으로 잇는다 — 노드 안(H100은 8개, B200 NVL72는 72개까지)은 직접 연결되고 그보다 큰 토폴로지는 GPU마다 O(log N)번의 홉을 거친다. GPU는 그래서 적은 홉으로 임의의 데이터를 보낼 수 있지만, TPU는 NVLink 스위치 같은 비싼 장비 없이 배선이 단순하고 장치당 링크 수·대역폭이 일정해 훨씬 큰 토폴로지로 늘어난다.

5. 대역폭 넷 — HBM·ICI·PCIe·DCN

TPU v5p의 대역폭 넷 (칩 하나 기준, 공표치)

PCIe와 DCN은 세대마다도 갈린다 — PCIe는 v6e가 초당 `3.2e10`바이트로 v4의 두 배이고, DCN은 v6e·TPU7x가 초당 `1.25e10`바이트, v5e가 초당 `3.125e9`바이트로 더 낮다. ICI는 DCN보다는 훨씬 빠르지만 HBM 대역폭에는 못 미친다. 모델을 여러 칩에 나눠 쪼갤 때는 이 순서(HBM보다 느린 ICI, ICI보다 느린 DCN)를 넘어서는 통신이 MXU를 굶기지 않도록 조심해야 한다. ICI로 이은 TPU 묶음을 저자들은 슬라이스(slice)라고 부르는데, 서로 다른 슬라이스는 DCN으로 잇는다. DCN은 호스트와 호스트를 잇는 연결이라, 슬라이스 사이에서 버퍼를 옮기려면 PCIe로 호스트까지 나갔다가 네트워크를 타고 상대 호스트로 들어간 뒤 다시 PCIe로 HBM까지 들어가야 한다.

저자들이 꼽는 핵심 요지 둘. ① 가중치 행렬은 두 축 모두 최소 128(v6e는 256)까지 채워야 MXU를 온전히 쓴다 — 더 작은 축은 128까지 패딩된다. ② 저정밀도 행렬곱이 대체로 더 빠르다. int8이나 int4는 세대에 따라 bf16보다 2배·4배 빠르지만 VPU 연산은 여전히 fp32로 돈다.

6. 스펙표 — 세대마다 얼마나 다른가

TPU 세대별 스펙 (칩 하나 기준, 공표치)

호스트 크기는 CPU 호스트 하나에 붙은 TPU의 토폴로지다. v5e는 호스트 하나에 4×2로 TPU 8개가 붙는다. 표에서 눈에 띄는 것은 세대가 오를수록 HBM 대역폭과 FLOPs/s가 나란히 오르지만은 않는다는 점이다. v5e는 추론에 맞춰 HBM 용량을 줄인 대신 포드를 16×16 2차원 토러스로 넓게 늘렸고, v5p는 HBM 용량을 키우고 3차원 토러스로 학습에 맞췄다. TPU7x는 이 표 안에서 칩당 HBM 용량과 bf16 FLOPs/s 둘 다 가장 크다.

7. 숫자로 풀어보기 — 두 문제

저자들이 낸 연습문제 중 둘이 이 장의 숫자를 실제로 어떻게 쓰는지 보여준다.

문제 1은 200B 파라미터 모델을 bf16으로 TPU v4p 32개에 나눠 놓고 샘플링할 때 HBM에서 시스톨릭 배열로 파라미터를 전부 읽어 오는 데 얼마나 걸리는지 묻는다. 옮길 바이트는 `sizeof(bf16) * 200e9 = 400e9`, 칩 32개로 나누면 칩당 `12.5e9`바이트이고 칩당 HBM 대역폭이 `1.23e12`바이트/초이니 적재에 약 10밀리초가 걸린다. 저자들은 이 10ms를 그 모델을 샘플링하는 지연시간의 현실적인 하한으로 본다 — 샘플링 한 걸음마다 파라미터를 HBM에서 전부 읽어야 하니 배치 크기가 작을 때는 실제 속도가 이 하한에 가깝게 붙는다.

문제 6은 여러 대역폭을 한 번에 엮는다. int8 행렬 A(128×1,024 제곱, 약 16GB)를 TPU v5e 4×4 슬라이스에 고르게 나눠 각 칩의 호스트 DRAM에 얹어 두고, 이를 전부 TPU{0,0}로 모아 bf16 벡터와 곱하는 문제다. 저자들이 짚은 경로는 넷이다. ① PCIe로 각 호스트가 가진 조각(8GB씩, 링크 16개)을 자기 TPU로 올리는 데 약 63ms가 걸린다. ② ICI로 TPU{0,0}가 나머지 15GB를 두 축·링크당 `4.5e10`바이트/초로 받아오는 데 약 167ms가 걸린다. ③ HBM에서 MXU로 `16e9`바이트를 옮기는 데 약 20ms가 걸린다. ④ 실제 곱셈(`2.7e11` FLOPs, `1.97e14` FLOPs/s)에는 약 1.4ms가 걸린다. TPU가 이 넷을 최대한 겹쳐 돌린다고 보면 전체 시간은 가장 느린 조각인 ICI 복사, 167ms에 가깝게 잡힌다 — 겹침이 완벽하지 않으면 200ms에 가까워진다.