Google TPU

Google Cloud TPUs는 torch_tpu 패키지를 통해 PyTorch를 네이티브로 실행하며, 이 패키지는 tpu 장치를 등록합니다. torch_tpu는 확인하는 각 입력 형식에 대해 실행할 연산을 컴파일하므로, 입력 형식을 정적으로 유지하면 재컴파일을 방지할 수 있습니다.

설치

torch_tpu 저장소의 지침에 따라 torch_tpu를 설치합니다. torch를 가져오면 torch_tpu가 자동으로 로드됩니다. PyTorch가 TPU 칩을 인식하는지 확인합니다.

import torch

print(torch.accelerator.is_available(), torch.accelerator.device_count())

True와 TPU 칩의 수를 출력합니다. 예를 들어 True 4입니다.

추론 실행

device_map="auto"는 TPU를 감지하지 않습니다. 자동 배치에는 사용 가능한 장치를 나열하는 Accelerate에 의존하며, Accelerate는 TPU 장치를 나열하지 않으므로 모델이 CPU에 조용히 로드됩니다. from_pretrained()에 device_map="tpu"를 명시적으로 전달하거나, Pipeline에 device="tpu"를 전달하세요.

.to(model.device)를 호출하는 대신 생성 입력을 CPU로 전달합니다. generate()는 forward가 소비하는 텐서만 TPU로 이동합니다. 각 단계에서 커지는 텐서(출력 시퀀스, attention_mask)는 CPU에 유지되므로 컴파일된 그래프가 안정적으로 유지됩니다. generate는 입력과 동일한 디바이스에서 출력을 반환합니다. 자세한 내용은 생성 텐서를 CPU에 유지하기를 참조하세요. TPU에서는 모델이 기본적으로 SDPA를 사용합니다.

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B", dtype="bfloat16", device_map="tpu"
)

# leave the inputs on CPU
inputs = tokenizer("The French Bread Law states", return_tensors="pt")
output = model.generate(
    **inputs, do_sample=True, max_new_tokens=50, cache_implementation="static"
)
print(tokenizer.decode(output[0], skip_special_tokens=True))

The French Bread Law states that a farmer can sell their harvest of bread at a certain price. The price is determined by the market price and the price of the bread, and the price is set by the farmers' decision-making committee. What are the implications of this law on

정적 캐시를 사용하는 경우, generate()는 torch.compile과 tpu 백엔드를 자동으로 사용하여 forward pass를 컴파일합니다. generate에 disable_compile=True를 전달하여 이 기능을 비활성화할 수 있습니다.

Tensor 병렬성

torch_tpu는 TPU 칩마다 프로세스 하나를 실행하며, 각 프로세스는 자신의 칩을 tpu:0으로 인식합니다. DistributedConfig를 텐서 병렬성 크기와 함께 생성하고 from_pretrained()에 전달합니다. Transformers는 로드 시점에 모델을 샤딩합니다.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DistributedConfig

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B")
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-8B",
    dtype=torch.bfloat16,
    distributed_config=DistributedConfig(tp_size=4),
)

# leave the inputs on CPU
inputs = tokenizer("The French Bread Law states", return_tensors="pt")
output = model.generate(
    **inputs, do_sample=False, max_new_tokens=50, cache_implementation="static"
)
if torch.distributed.get_rank() == 0:
    print(tokenizer.decode(output[0], skip_special_tokens=True))

torch_tpu는 프로세스가 시작되기 전에 몇 가지 환경 변수가 설정되어야 하며, torchrun은 이를 자체적으로 수행하지 않습니다. torch.tpu.distributed.set_tpu_launch_env()를 호출하고 Python에서 torchrun을 시작하세요.

# launch.py
import sys

import torch
from torch.distributed.run import main

torch.tpu.distributed.set_tpu_launch_env()
main(sys.argv[1:])

torchrun의 일반 인수를 런처에 전달합니다.

python launch.py --nproc_per_node=4 generate.py

일부 모델은 TPU 칩에 걸쳐 샤딩하기 전에 적응 작업이 필요합니다.

tp_size로 나누어 떨어지지 않는 key-value head 수를 가진 모델은 attention 레이어를 샤딩할 수 없으며, forward pass에서 shape 오류로 실패합니다. 대신 tp_plan에서 해당 규칙을 오버라이드하여 모든 칩에 attention 투영을 복제하세요. 오버라이드하지 않은 규칙은 모델의 기본 플랜을 유지합니다. 예를 들어, GLM-4-9B-0414는 2개의 key-value head를 가지며, 이는 4개의 칩에 나누어 배치할 수 없습니다.

model = AutoModelForCausalLM.from_pretrained(
    "zai-org/GLM-4-9B-0414",
    dtype=torch.bfloat16,
    distributed_config=DistributedConfig(
        tp_size=4,
        tp_plan={
            "layers.*.self_attn.q_proj": "replicated_with_grad_allreduce",
            "layers.*.self_attn.k_proj": "replicated_with_grad_allreduce",
            "layers.*.self_attn.v_proj": "replicated_with_grad_allreduce",
            "layers.*.self_attn.o_proj": "replicated_with_grad_allreduce",
        },
    ),
)

Qwen3-0.6B와 같이 가중치가 연결된 모델은 tp_size가 1보다 큰 경우 로드되지 않습니다.

Hub 커널 사용

kernels 라이브러리는 v0.17부터 TPU를 지원합니다. TPU 커널은 JAX와 Pallas 코드를 래핑하므로 kernels와 함께 jax를 설치하세요.

pip install "transformers[kernels]" jax

Hub에서 커널을 로드하는 방법은 커널 로드를 참조하세요.

한계

device_map="auto"는 TPU 대신 CPU에 모델을 로드합니다. device_map="tpu"를 명시적으로 전달하세요.
동적 shape는 재컴파일을 유발합니다. 입력 shape가 단계별로 변경되면 생성 속도가 느려집니다. 더 빠른 생성을 위해 정적 캐시를 사용하세요.
각 프로세스는 TPU 칩 하나만 액세스할 수 있습니다. 여러 칩에서 실행하려면 칩당 프로세스 하나를 사용하세요.
TPU에서 학습 및 양자화를 위한 Transformers 통합은 아직 진행 중이므로 일부 기능이 아직 작동하지 않을 수 있습니다. 문제가 발생하면 이슈를 열어 주세요.

Update on GitHub