TPU microbenchmark 활용법: LLM 학습과 추론 병목을 추측하지 않고 찾는 절차

요약: Google의 TPU microbenchmark suite는 Network, Compute, HBM, Host Transfer, Attention 계층을 따로 측정해 병목을 분리합니다. 개발팀이 얻을 수 있는 핵심은 “GPU/TPU 스펙표”가 아니라 workload별 roofline 기준, MFU, TTFT, TPOT, communication stall을 실제 숫자로 확인하는 운영 습관입니다.

왜 microbenchmark가 필요한가

대규모 모델 학습이나 추론 인프라를 다루는 팀은 자주 비슷한 대화를 합니다. “TPU를 더 붙이면 빨라질까?”, “batch size를 늘리면 될까?”, “attention kernel을 바꿔야 하나?”, “네트워크가 병목인가?” 문제는 대부분 추측으로 시작된다는 점입니다.

하드웨어 스펙표는 peak 성능을 보여주지만, 실제 workload가 그 성능을 얼마나 쓰는지는 별개의 문제입니다. 모델은 compute만 하지 않습니다. HBM에서 데이터를 읽고, host와 device 사이를 오가고, multi-chip collective communication을 하고, attention primitive를 수행합니다. 어느 하나가 막히면 전체 처리량이 떨어집니다.

Google이 공개한 TPU microbenchmark 글은 이런 추측을 줄이는 방법을 설명합니다. 핵심은 전체 모델을 먼저 돌려보는 것이 아니라, 하드웨어 구성 요소별 baseline을 측정하고, 실제 workload가 어디에 걸리는지 roofline model로 분류하는 것입니다.

무엇을 측정해야 하나

Google의 microbenchmark suite는 다섯 영역을 봅니다.

첫째, Network입니다. multi-chip 학습에서 all-gather, all-reduce, reduce-scatter, all-to-all 같은 collective operation은 필수입니다. 여기서는 throughput(GB/s)과 latency를 봅니다. 모델이 커질수록 compute보다 communication이 병목이 되는 일이 흔합니다.

둘째, Compute입니다. GEMM/Matmul을 중심으로 raw computational throughput을 측정합니다. 여기서 중요한 지표는 TFLOPs와 MFU(Model FLOPs Utilization)입니다. MFU가 낮다면 accelerator의 matrix unit을 제대로 쓰지 못하고 있다는 뜻입니다.

셋째, HBM입니다. high bandwidth memory에서 데이터를 얼마나 빨리 읽고 쓰는지 봅니다. attention, routing, sparse activation처럼 memory access가 많은 workload는 HBM bandwidth가 실제 병목이 될 수 있습니다.

넷째, Host Transfer입니다. CPU와 accelerator HBM 사이의 H2D, D2H 전송률을 봅니다. 데이터 로딩이나 preprocessing이 host-device transfer에 막히면 accelerator가 놀게 됩니다.

다섯째, Attention 관련 primitive입니다. Google 글에서는 Ragged-Paged Attention, tokamax splash attention, BMM 같은 transformer inference 핵심 primitive를 따로 측정합니다. 이는 TTFT(Time-to-First-Token)와 TPOT(Time-Per-Output-Token)을 예측하는 데 중요합니다.

roofline으로 병목을 분류하는 법

microbenchmark의 목적은 “숫자를 많이 모으는 것”이 아닙니다. 목적은 workload를 compute-bound, memory-bound, network-bound로 분류하는 것입니다.

compute-bound라면 MXU plateau에 걸린 상태입니다. 이 경우 HBM 접근을 조금 줄여도 큰 효과가 없습니다. kernel efficiency, FLOPs 감소, 연산 shape 정렬이 중요합니다.

memory-bound라면 HBM bandwidth가 한계입니다. data locality, activation checkpointing, rematerialization 전략, attention 구현 변경이 효과적일 수 있습니다.

network-bound라면 ICI 또는 DCN communication이 accelerator를 기다리게 만듭니다. sharding 전략, FSDP 설정, communication overlap, topology 선택을 봐야 합니다.

이 분류 없이 튜닝하면 위험합니다. network-bound 문제에 kernel을 바꿔도 효과가 없고, memory-bound 문제에 chip을 더 붙이면 communication만 늘 수 있습니다.

사례에서 볼 점: 110B MoE workload

Google 글의 사례는 4x4x4 TPU 7x 구성에서 110B Mixture-of-Experts 학습 workload를 분석합니다. 결과는 한 가지 병목이 아니라 구간별 병목이었습니다.

forward pass와 dense-core operation은 compute-bound였고 1.85 PFLOPS에 도달했습니다. 반면 routing과 attention primitive는 HBM saturation 때문에 memory-bound였으며 SOL 대비 30~60% 수준에 머물렀습니다. sparse activation pattern에서 expert dispatch가 I/O-heavy해졌기 때문입니다.

microbenchmark 결과는 실제 intervention으로 이어졌습니다. ICI collective benchmark는 communication stall을 찾는 데 쓰였고, HBM bandwidth test는 memory-bound attention primitive를 확인했습니다. 그 결과 SparseCore collective offloading과 Tokamax Splash Attention 같은 선택이 나왔고, training step time이 21.2% 줄었다고 설명합니다.

이 사례의 교훈은 “큰 모델은 항상 compute-bound”가 아니라는 것입니다. MoE, retrieval, tool-using model, long context inference는 각 단계마다 병목이 달라질 수 있습니다.

실무 적용 순서

첫째, full workload benchmark 전에 hardware baseline을 잡습니다. 같은 TPU slice라도 topology, region, runtime, driver, container 설정에 따라 실제 숫자가 달라질 수 있습니다.

둘째, microbenchmark 결과를 저장합니다. 일회성 측정이 아니라 regression 기준이 되어야 합니다. 새 이미지, 새 JAX 버전, 새 kernel을 넣었을 때 baseline이 흔들리는지 봐야 합니다.

셋째, end-to-end trace와 연결합니다. microbenchmark에서 HBM bandwidth가 낮게 나왔다고 바로 원인이 확정되는 것은 아닙니다. 실제 workload trace에서 stall 위치와 함께 봐야 합니다.

넷째, 튜닝 후보를 병목별로 나눕니다.

• compute-bound: kernel selection, shape alignment, FLOPs reduction.
• memory-bound: attention kernel, rematerialization, activation layout.
• network-bound: sharding, topology, collective overlap.
• host-bound: data pipeline, prefetch, H2D transfer.

다섯째, 비용 기준을 함께 봅니다. step time이 줄어도 accelerator 수가 늘어 비용이 증가하면 제품 관점에서는 손해일 수 있습니다.

체크리스트

• peak TFLOPs만 보고 인프라 결정을 내리지 않는다.
• Network, Compute, HBM, Host Transfer, Attention을 분리해 측정한다.
• MFU, TTFT, TPOT, step time, GB/s를 함께 기록한다.
• roofline으로 compute-bound, memory-bound, network-bound를 먼저 분류한다.
• full training run 전에 작은 slice에서 baseline을 만든다.
• 튜닝 전후 결과를 같은 config, 같은 데이터, 같은 topology에서 비교한다.
• microbenchmark 결과를 CI 또는 정기 regression check에 연결한다.

LLM 인프라 최적화에서 가장 비싼 실수는 틀린 병목을 고치는 것입니다. microbenchmark는 화려한 도구가 아니라, 그 실수를 줄이는 계측 습관입니다.