TPU vs GPU: 구글 클라우드에서 LLM 훈련을 위한 실전 성능 테스트

Source: Dev.to

죄송합니다만, 현재 저는 외부 웹사이트의 내용을 직접 가져올 수 없습니다. 번역이 필요한 텍스트를 그대로 제공해 주시면, 해당 부분을 한국어로 번역해 드리겠습니다.

소개

대형 언어 모델(LLM)이 규모를 계속 확대함에 따라, 학습에 사용되는 기본 하드웨어는 프로젝트 성공에 있어 가장 중요한 요소가 되었습니다. 현재 업계는 흥미로운 아키텍처 전쟁에 휩싸여 있습니다: NVIDIA GPU의 범용 성능 대 Google TPU(텐서 처리 장치)의 목적에 맞춘 효율성.

Google Cloud Platform(GCP)에서 엔지니어와 아키텍트가 A100/H100 GPU 클러스터와 TPU v4/v5p 팟 중 하나를 선택할 때, 이는 단순히 비용 문제를 넘어 소프트웨어 아키텍처, 데이터 파이프라인, 그리고 수렴 속도에 직접적인 영향을 미칩니다. 이 글에서는 실제 LLM 학습 성능을 기준으로 두 아키텍처를 심층 기술 분석합니다.

실리콘 수준 차이

근본적인 차이는 칩이 행렬 곱셈을 처리하는 방식에 있으며, 이는 Transformer 아키텍처의 핵심 연산입니다.

클러스터‑레벨 통신

LLM(예: Llama‑3 또는 GPT‑4) 훈련은 절대 단일 칩에서 이루어지지 않으며, 클러스터에서 수행됩니다. 칩 간 통신 속도가 종종 순수 TFLOPS보다 더 큰 영향을 미칩니다.

• NVIDIA

  • NVLink/NVSwitch: 노드 내부 통신.
  • InfiniBand: 노드 간 통신.
  • H100은 NVLink 4를 지원하며 ≈ 900 GB/s 대역폭을 제공합니다.

• Google TPUs

  • **Optical Circuit Switch (OCS)**와 독점적인 Inter‑Core Interconnect (ICI).
  • TPU v4 및 v5p는 OCS를 활용해 팟 토폴로지를 동적으로 재구성하고, 수천 개 칩에 걸쳐 전통적인 네트워킹 계층의 오버헤드 없이 저지연·고대역폭 통신을 제공하는 거대한 3‑D 토러스를 형성합니다.

기능 비교

Real‑World Test Setup

우리는 Google Cloud에서 7 B‑parameter Transformer 모델(Llama‑2 아키텍처)의 학습을 수행했습니다.

Example JAX Sharding Code (runs on both TPU pods and multi‑GPU setups)

import jax
import jax.numpy as jnp
from jax.sharding import Mesh, PartitionSpec, NamedSharding
from jax.experimental import mesh_utils

# Detect devices (TPU or GPU)
devices = jax.devices()
print(f"Devices found: {devices}")

# Define a 2‑D mesh for model and data parallelism
# Works identically on TPU pods and multi‑GPU setups
device_mesh = mesh_utils.create_device_mesh((4, 2))
mesh = Mesh(devices=device_mesh, axis_names=('data', 'model'))

# Create a sharded array
# 'data' axis shards the batch, 'model' axis shards the weights
sharding = NamedSharding(mesh, PartitionSpec('data', 'model'))

def train_step(state, batch):
    # XLA handles the communication primitives (all‑reduce)
    # during the gradient computation automatically
    def loss_fn(params):
        logits = model.apply(params, batch['input'])
        return jnp.mean(cross_entropy(logits, batch['target']))

    grads = jax.grad(loss_fn)(state.params)
    return state.apply_gradients(grads=grads)

# JIT‑compile the step for XLA optimization
parallel_train_step = jax.jit(train_step)

Performance Results

TPU의 결정론적 실행과 ICI 인터커넥트가 대기 시간을 최소화하여, 칩당 원시 처리량이 약간 낮음에도 전체 활용도가 높아집니다.

분산 학습 전략

비용 고려 사항

성능은 비용과 비교해서 평가해야 합니다:

• Google Cloud TPU 가격은 동등한 컴퓨팅 시간에 대해 일반적으로 H100 가격보다 낮습니다.
• Spot TPU는 온‑디맨드 인스턴스보다 70 %까지 저렴할 수 있습니다.
• GPU도 Spot 인스턴스를 제공하지만, 가격 차이와 가용성은 지역 및 수요에 따라 다릅니다.

요약

1. Raw throughput vs. utilization – H100은 소규모 배치와 칩당 속도에서 앞서며, TPU는 대규모에서 지속적인 활용도에서 뛰어납니다.
2. Interconnect matters – TPU의 광 회로 스위치는 수천 개 칩에 걸쳐 보다 원활하게 확장되는 토폴로지를 제공합니다.
3. Software ecosystem – 두 플랫폼 모두 이제 XLA를 지원합니다; JAX는 TPU에 자연스럽게 맞으며, GPU에서는 NCCL을 사용하는 PyTorch가 표준입니다.
4. Cost efficiency – Spot TPU는 종종 GCP에서 대규모 LLM 훈련에 가장 좋은 가격 대비 성능 비율을 제공합니다.

올바른 하드웨어 선택은 궁극적으로 워크로드의 배치 크기, 원하는 훈련 속도, 예산 제약에 달려 있습니다. 위에서 설명한 아키텍처의 미묘한 차이를 이해함으로써 성능 목표와 비용 목표 모두에 부합하는 현명한 결정을 내릴 수 있습니다.

가용성 및 비용 비교

대규모 연속 블록의 H100 GPU 가용성은 TPU 슬라이스보다 종종 낮습니다.

예시 비용 비교 (8‑칩 노드에 대한 예상 시간당 비용)

Tokens per Dollar를 계산할 때, TPU v5p는 H100이 약간 높은 원시 처리량을 가지고 있음에도 불구하고 우리 훈련 실행에서 15–25 % 더 우수한 성능을 보였습니다. 이는 예산이 주요 제약인 장기 사전 학습 단계에서 TPUs가 선호되는 선택이 됨을 의미합니다.

GPU가 여전히 빛나는 경우

• Ecosystem & Flexibility – 대부분의 오픈‑소스 ML 연구는 먼저 CUDA용으로 작성됩니다. 틈새 라이브러리나 최신 어텐션 메커니즘(예: FlashAttention‑3)은 보통 NVIDIA를 먼저 최적화합니다.
• Torch‑XLA는 PyTorch가 TPU에서 실행되도록 하지만, CPU와 TPU 간의 “컨텍스트 스위칭”을 피하기 위해 종종 약간의 코드 수정을 필요로 하며, 이는 성능을 크게 저하시킬 수 있습니다.
• Debugging – XLA 코드는 컴파일되기 때문에 학습 루프 안에 print 문을 바로 넣을 수 없습니다. jax.debug.print 또는 Cloud TPU 프로파일러를 사용해 HBM stalls나 Infeed queues와 같은 병목 현상을 식별하세요.

일반적인 TPU 병목 현상: Infeed

TPU를 사용할 때 자주 발생하는 제한은 Infeed이며, CPU가 데이터를 충분히 빠르게 공급하지 못해 TPU가 대기하게 됩니다.

# Using the TPU Profiler in a training loop
import jax

with jax.profiler.trace("/tmp/tpu_profile", create_perfetto_link=True):
    for i in range(100):
        state = parallel_train_step(state, next(data_iter))
        # Ensure the TPU doesn't wait for the host
        if i % 10 == 0:
            print(f"Step {i} completed")

Google Cloud에서 LLM 훈련을 위한 의사결정 트리

“TPU vs GPU” 논쟁은 이제 순수 속도가 아니라, 여러분의 특정 워크로드에 대한 시스템 수준 효율성에 관한 것입니다.

강점 요약

모델 아키텍처와 예산을 신중히 평가함으로써, LLM 훈련 프로젝트를 일정과 예산 내에 유지할 수 있는 적절한 가속기를 선택할 수 있습니다.

추가 읽을거리 및 자료

• Technical Guides – Google Cloud AI Architecture & Implementation
• 팔로우하기:
  • Twitter / X
  • LinkedIn
  • GitHub