저자들은 H100의 bf16 처리량 9.89e14 FLOPs/s와 TPU v6e의 9.1e14 FLOPs/s를 놓고 계산 시간(T_math)과 통신 시간(T_comms) 중 큰 값이 실행 시간의 하한이라는 루프라인 모델을 세운다. TPU v5e MXU의 임계 산술강도는 1.97e14 FLOPs/s를 8.2e11 bytes/s로 나눈 240 FLOPs/byte이고, bf16 행렬곱은 근사를 거쳐 배치 크기가 240 토큰을 넘어야 연산 병목에 들어선다는 규칙으로 정리된다. 두 벡터의 내적은 산술강도가 N이 커져도 1/2에 그쳐 거의 항상 통신에 묶이는 사례로, 두 칩이 절반씩 곱한 행렬을 다시 합칠 때는 D가 8,755를 넘어야 연산이 통신을 앞선다는 사례로 함께 다룬다. 가중치만 int8로 낮추고 연산은 bf16을 유지하면 임계 배치가 240에서 120으로 절반이 된다는 계산도 실려 있다.
한줄 코멘트. 이 장의 결론은 간단하다. 가속기가 노는 시간은 대개 연산기가 모자라서가 아니라 데이터를 옮기는 데 시간을 다 써서 생긴다. 저자들은 계산 시간과 통신 시간 중 큰 값이 실행 시간을 정한다는 루프라인(roofline) 모델을 세우고, 산술강도라는 잣대 하나로 어느 쪽이 병목인지 미리 계산해 낸다. TPU v5e에서 이 잣대가 240 FLOPs/byte라는 구체적인 숫자로 떨어지는 과정을 보여준다.
저자들은 이 장을 단순한 물음 하나로 연다. 어떤 알고리즘이 왜 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는 가격이 다르므로 처리 속도만으로 어느 쪽이 유리한지 가릴 수는 없다.
칩 안에서는 텐서를 고대역폭 메모리(HBM — 칩 옆에 쌓은 주 메모리)와 연산 코어 사이로 옮겨야 한다. H100은 이 대역폭이 초당 3.35테라바이트, TPU v6e는 초당 1.6테라바이트다. 모델을 여러 가속기에 나눠 실으면 텐서는 칩과 칩 사이도 오간다. 이때 쓸 수 있는 경로는 보통 ICI(칩 사이를 잇는 전용 연결)·DCN(데이터센터 네트워크)·PCIe 셋이고, 셋 다 대역폭이 다르다.
칩 안이든 칩 사이든 통신 시간은 같은 식으로 구한다. 옮길 바이트 수를 대역폭(초당 바이트)으로 나눈다(`T_comms = 통신 바이트 / 대역폭 바이트/s`). 계산과 통신은 대개 겹쳐 돌릴 수 있으므로 실행 시간의 하한은 둘 중 큰 값이고, 전혀 안 겹친다고 보수적으로 잡으면 상한은 둘을 더한 값이다. 저자들은 실무에서 이 하한을 기준으로 최적화한다고 밝힌다. 겹쳐 쓰기를 최대한 살리면 실제 시간이 하한에 가까워지기 때문이다. 최댓값을 기준으로 잡으면 하한과 상한의 차이도 묶인다. 상한(둘의 합)은 하한(둘 중 큰 값)의 최대 두 배를 넘지 않는다. 이보다 더 정확한 값이 필요하면 계산과 통신이 겹치는 구간과 그 밖의 오버헤드까지 따져야 하는데, 이는 실제 모델을 프로파일링해야 나오는 값이라고 저자들은 덧붙인다.
시간의 하한과 상한은 어디서 나오나
계산이 통신보다 오래 걸리면 연산기를 놀리지 않고 다 쓴다는 뜻이라 "연산 병목(compute-bound)"이라 부른다. 반대로 통신이 더 오래 걸리면 연산기 일부가 데이터를 기다리며 논다는 뜻이라 "통신 병목(comms-bound)"이라 부른다. 이 둘을 나누는 잣대가 산술강도다.
산술강도(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를 대역폭으로 나눈 값(가속기의 산술강도)보다 크다는 뜻으로 바뀐다. 두 산술강도만 견주면 실제로 돌려 보지 않고도 어느 쪽이 병목인지 미리 알 수 있는 이유가 여기에 있다.
산술강도 240이 나누는 것
산술강도가 TPU v5e의 임계값(240 FLOPs/byte) 위냐 아래냐
이 잣대가 얼마나 낮게 나올 수 있는지 보여주는 예가 두 벡터의 내적이다. 길이 N짜리 bf16 벡터 두 개를 곱해 더하려면 2N바이트씩 두 번 읽고 N번 곱하고 N-1번 더한 뒤 2바이트를 다시 써야 한다. N이 커질수록 이 비율은 1/2로 수렴한다. 240은커녕 1에도 못 미치는 값이라 내적은 어떤 하드웨어에서도 거의 항상 통신 병목에 걸린다고 저자들은 말한다. 다만 저자들은 각주에서 내적이 실제로는 행렬곱 유닛(MXU)이 아니라 벡터 처리 유닛(VPU)에서 돈다는 점을 짚는다. TPU v5p의 VPU는 코어당 초당 약 7e12 FLOPs를 처리해 임계 산술강도가 약 3에 불과하다. 산술강도 1/2인 내적은 이 낮은 기준으로 봐도 여전히 통신 병목에 걸린다.
이 관계를 그래프 하나로 정리한 것이 루프라인 도표다. 가로축에 산술강도, 세로축에 실제로 달성하는 FLOPs/s를 놓는다. 산술강도가 낮은 구간에서는 그래프가 직선으로 올라간다. 대역폭이 한계라 산술강도가 늘어난 만큼 처리량도 그대로 늘어나기 때문이다. 임계 산술강도를 넘어서면 그래프는 수평으로 꺾인다. 이미 연산기를 다 쓰고 있어서 산술강도를 더 올려도 처리량이 늘지 않는다. 저자들이 든 예시 그래프에서는 서로 다른 산술강도를 가진 두 알고리즘(Algo 1·Algo 2)을 두 가지 대역폭(BW1·BW2) 아래 나란히 놓는다. 대역폭을 BW1에서 BW2로 키우면 꺾이는 지점이 오른쪽 위로 옮겨 가면서 같은 산술강도에서도 더 높은 처리량을 낼 수 있다. 이 그래프는 구간을 셋으로 나눈다. 두 대역폭 모두에서 통신에 발목 잡히는 구간, 낮은 대역폭 BW1에서만 병목이 걸리는 구간, 어느 대역폭에서도 연산기를 다 쓰는 구간이다. 가운데 구간에 있는 알고리즘은 대역폭만 BW1에서 BW2로 올려도 병목에서 벗어난다.
행렬곱 `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토큰에 그친다. 여기서 배치는 시퀀스 개수가 아니라 토큰 개수로 잰다는 점도 저자들은 강조한다. 시퀀스가 같든 다르든 토큰 수만 맞으면 루프라인 계산은 동일하다. 이 임계 배치는 모델을 여러 칩에 나눠 싣더라도(샤딩) 그대로 적용되는데, 샤딩으로 칩 수를 늘리면 연산량과 대역폭이 함께 늘어나므로 임계 배치는 가중치 한 벌을 기준으로 정해지기 때문이라고 저자들은 짚는다.
실제로 큰 행렬곱을 실행할 때는 행렬 전체를 한 번에 못 올리고 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과 크게 다르지 않다고 저자들은 밝힌다.
가속기 두 개에 행렬을 D 차원 기준으로 절반씩 나눠 곱할 때도 같은 계산이 적용된다. 각 칩은 전체 계산량의 절반만 맡으므로 계산 시간은 그대로 절반이 된다. 두 칩은 각자 구한 부분합을 상대 칩으로 보내 더해야 하므로, 통신 바이트는 2BF이고 통신 시간은 이를 칩 간 대역폭(저자들이 예로 든 초당 4.5e10바이트)으로 나눈 값이다.
칩 두 개로 나눈 행렬곱, 통신이 끼어드는 자리
이 경우 연산 병목이 되는 조건은 `D/2 > 4377`, 즉 `D > 8755`다. 앞서 배치 크기 B에 걸려 있던 조건이 이번에는 D에 걸린다는 점이 다르다. 저자들은 이 예시가 다소 인위적으로 고른 것이라고 인정하면서도, 이 책이 다루는 루프라인 대부분은 칩 안이 아니라 칩 사이 통신 쪽이라고 못박는다. 여러 TPU에 걸쳐 나뉜 행렬곱을 언제 병렬로 돌려도 되는지 판단하는 데 바로 이런 계산이 결정적이라는 것이다. 칩 안 통신이든 칩 사이 통신이든 같은 방식으로 하한과 상한을 구할 수 있다는 것이 이 장의 요지라고 저자들은 정리한다.