PyTorch/XLA에서 TorchTPU로 옮기기: 11월 15일 아카이브 전 코드 변경과 검증 순서

PyTorch/XLA 저장소 README에는 이 프로젝트가 TorchTPU로 대체됐으며 2026년 11월 15일에 저장소가 아카이브된다는 공지가 올라와 있다. Cloud TPU에서 torch_xla로 학습이나 추론을 돌리고 있다면 아카이브 이후에는 수정 릴리스를 기대하기 어렵다. 이 글은 기존 torch_xla 코드를 TorchTPU 0.2.0 기준으로 옮기는 설치, 코드 변경, 실행 모드 선택, 결과 검증 순서를 다룬다.

PyTorch/XLA와 TorchTPU는 실행 방식부터 다르다

2026년 10월 20일 공개된 PyTorch 블로그의 TorchTPU 소개 글은 PyTorch/XLA의 한계를 세 가지로 짚는다. Lazy Tensor 기반이라 컴파일 영역을 코드에서 명시해야 했고, 컴파일 지연이 컸으며, PyTorch의 기본 실행 방식인 eager 모드를 지원하지 못했다는 것이다. TorchTPU는 PyTorch 디스패처에 PrivateUse1 백엔드로 등록되는 네이티브 백엔드이며, 런타임에 JAX 의존성이 없고 torch.compile도 TorchDynamo 위에서 동작한다. Google이 Meta의 PyTorch 팀과 함께 개발하며 Apache 2.0 라이선스로 공개됐다.

코드 수준에서 가장 큰 변화는 단계 경계를 직접 표시하던 코드가 사라진다는 점이다. Strict 모드는 블로그 설명대로 GPU에서 PyTorch를 eager로 돌릴 때와 같은 방식으로 동작하므로, 결과를 CPU로 가져오는 .item()·.cpu() 시점에만 TPU 연산 완료를 기다린다.

옮기기 전에 맞춰야 할 환경 조건

현재 torch_xla 작업이 v4 이하 세대에서 돌고 있다면 FAQ의 지원 목록에 들어 있지 않으므로, 코드보다 하드웨어 세대 전환을 먼저 결정해야 한다. 또 PyTorch/XLA README의 설치 예시는 torch==2.8.0과 torch_xla[tpu]==2.8.0을 짝으로 고정하므로, 기존 환경에 TorchTPU를 덧설치하면 torch 버전 요구가 엇갈릴 수 있다. 새 가상환경이나 새 컨테이너 이미지에서 시작하는 것이 이전 전후를 비교하기에도 깔끔하다. 이미지를 학습용과 추론용으로 나누는 기준은 Docker로 AI 애플리케이션 환경 만들기에서 다룬 방식을 그대로 쓸 수 있다.

설치하고 TPU 디바이스가 잡히는지 확인하기

1. 새 가상환경 만들기 — TPU VM에서 Python 3.12 가상환경을 만든다. 기존 torch_xla 환경은 비교용으로 남겨 둔다.
2. torch_tpu 설치 — 버전을 0.2.0으로 고정하고 CPU용 PyTorch 인덱스를 함께 지정한다. 저장소 README에 따르면 이 인덱스를 빼면 pip가 TPU에 필요 없는 CUDA 빌드를 받는다.
3. 인식 확인 스크립트 실행 — 텐서를 "tpu" 디바이스에 만들고 연산 결과를 CPU로 가져와 본다.

sudo apt update && sudo apt install -y python3.12 python3.12-venv
python3.12 -m venv env
source env/bin/activate
pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu

# check_tpu.py
import torch

x = torch.randn((4, 4), device="tpu")
y = torch.ones((4, 4), device="tpu")
result = torch.matmul(x, y).cpu()
print(x.device, result.shape)

import torch_tpu가 없어도 되는 이유는 TorchTPU가 PyTorch의 autoload 메커니즘으로 등록되기 때문이다. 출력에 tpu 디바이스가 찍히고 torch.Size([4, 4])가 나오면 설치는 끝난 것이다. 여기서 디바이스 오류가 나면 TPU VM이 아닌 곳에서 실행했거나, CUDA 빌드 torch가 먼저 깔린 환경일 가능성이 높으므로 pip show torch로 빌드 종류부터 확인한다.

torch_xla 코드를 TorchTPU 코드로 바꾸는 대응표

# 이전: PyTorch/XLA
import torch
import torch_xla

model = MyModel().to("xla")
for inputs, labels in loader:
    with torch_xla.step():
        inputs, labels = inputs.to("xla"), labels.to("xla")
        optimizer.zero_grad()
        loss = loss_fn(model(inputs), labels)
        loss.backward()
        optimizer.step()
torch_xla.sync()

# 이후: TorchTPU
import torch

device = torch.device("tpu")
model = MyModel().to(device)
for step, (inputs, labels) in enumerate(loader):
    inputs, labels = inputs.to(device), labels.to(device)
    optimizer.zero_grad()
    loss = loss_fn(model(inputs), labels)
    loss.backward()
    optimizer.step()
    if step % 100 == 0:
        print(step, loss.item())  # 여기서 TPU 결과를 기다린다

로그를 매 스텝 찍던 코드는 그대로 두면 스텝마다 CPU가 TPU를 기다리게 되므로, 위처럼 간격을 두고 .item()을 부르는 것이 좋다. GPU와 TPU에서 같은 스크립트를 돌려야 한다면 블로그가 제시한 대안대로 torch.accelerator.current_accelerator()로 디바이스를 얻고 dist.init_process_group()을 백엔드 인자 없이 호출하는 방식을 쓸 수 있다. 이렇게 하면 디바이스 이름이 코드에 박히지 않는다.

Strict·Debug·DeferAndFuse 중 어떤 모드로 시작할까

블로그는 DeferAndFuse가 Strict보다 50% 이상 빨라질 수 있고 융합으로 임시 저장 공간이 줄어 피크 HBM 사용량도 낮아진다고 설명한다. 같은 글의 FLUX 학습 사례에서는 eager 기준 스텝 시간 11.49초가 DeferAndFuse로 9.95초, 여기에 SimpleFSDP·Pallas 커널·torch.compile(backend="tpu")까지 단계적으로 더해 3.02초가 됐다. 한 모델의 수치이므로 그대로 기대하기보다는 "Strict로 맞는지 확인 → DeferAndFuse → torch.compile" 순서로 한 단계씩 바꾸며 스텝 시간을 직접 재는 근거로 쓰면 된다.

torch.compile을 켤 때 입력 길이를 버킷으로 맞추기

블로그에 따르면 TPU에서 torch.compile을 호출하면 현재는 dynamic=False가 자동으로 설정된다. 동적 shape를 다루는 bounded dynamism(SymInt 지원)은 기능 현황 문서에서 아직 실험 단계로 분류돼 있다. 따라서 배치마다 시퀀스 길이가 달라지면 새 shape마다 재컴파일이 일어난다. 길이를 몇 개의 버킷으로 고정하면 컴파일 횟수를 버킷 수만큼으로 제한할 수 있다.

import torch.nn.functional as F

BUCKETS = (128, 256, 512, 1024)

def pad_to_bucket(input_ids, attention_mask, pad_id):
    length = input_ids.size(1)
    target = next(b for b in BUCKETS if b >= length)
    extra = target - length
    input_ids = F.pad(input_ids, (0, extra), value=pad_id)
    attention_mask = F.pad(attention_mask, (0, extra), value=0)
    return input_ids, attention_mask

compiled = torch.compile(model, backend="tpu")

TorchTPU는 컴파일 결과를 프로세스 메모리(tier-1), 같은 호스트의 공유 메모리(tier-2), 원격 스토리지(tier-3) 세 계층에 캐시한다. tier-3는 기본으로 꺼져 있고 공유 스토리지 경로를 지정해야 켜지므로, 작업 재시작이 잦다면 API 레퍼런스에서 경로 설정 방법을 확인해 켜 두면 재시작 후 컴파일 시간을 줄일 수 있다.

GPU·CPU 결과와 숫자가 어긋나면 정밀도 설정부터 본다

CUDA에서 PyTorch의 기본 행렬곱 정밀도는 highest(순수 FP32)지만 TorchTPU의 기본값은 medium이다. medium은 입력을 BF16으로 바꿔 MXU에서 곱하고 누산은 FP32로 한다. 그래서 같은 모델이라도 GPU 결과와 소수점 아래 값이 다를 수 있으며, 이것만으로 이전이 잘못됐다고 판단하면 안 된다. 블로그에 따르면 torch.set_float32_matmul_precision("high")는 MXU 연산 비용이 약 3배, "highest"는 약 6배로 늘고 성능 경고가 한 번 출력된다. 수렴 문제가 확인된 특정 레이어만 높이려면 torch.tpu.precision 컨텍스트 매니저로 범위를 좁힌다.

차이가 허용 범위인지는 기준을 정해 놓고 판단한다. TorchTPU 팀은 비전·오디오·인코더 모델 검증에 상대 MSE(MSE를 기준 출력의 분산으로 나눈 값) 0.10 이하를 통과 기준으로 썼다. 같은 식으로 CPU 출력과 비교하는 스크립트는 다음과 같다.

import torch

torch.manual_seed(0)
model = MyModel().eval()
x = torch.randn(8, 512)

with torch.no_grad():
    ref = model(x)                       # CPU 기준값
    out = model.to("tpu")(x.to("tpu")).cpu()

rel_mse = (out - ref).pow(2).mean() / (ref.var() + 1e-7)
print(f"relative MSE: {rel_mse.item():.4e}")

LLM이라면 출력 텐서 비교보다 같은 평가셋의 perplexity나 골든셋 점수를 이전 전후로 비교하는 쪽이 판단하기 쉽다. 블로그의 LLM 검증 기준은 perplexity 상대 차이 0.10 이내였다. 평가셋을 코드로 고정해 두는 방법은 프로덕션 배포 전 LLM 평가 글의 회귀 비교 절차를 그대로 적용하면 된다.

분산 학습 스크립트를 옮길 때 막히는 지점

TorchTPU는 GPU의 PyTorch처럼 디바이스당 프로세스 하나를 쓰고, 이를 강제한다. 다만 TPU 프로세스는 자기 디바이스 하나만 볼 수 있어 호스트에 디바이스가 몇 개인지 모른다. 그래서 torchrun이 설정하지 않는 TORCH_TPU_TOPOLOGY(슬라이스 모양, 예: 2,2,1)와 TORCH_TPU_SLICEBUILDER_ADDRESSES(디바이스별 포트 목록)를 먼저 넣어야 한다. Google Cloud 환경에서는 아래처럼 torch.tpu.distributed 모듈이 값을 출력하게 해서 환경 변수로 만든다.

# 8칩 단일 호스트: 칩마다 프로세스 1개
eval $(python -m torch.tpu.distributed | sed 's/^/export /')
torchrun --nproc_per_node=8 train.py

• 프로세스 수는 칩이 아니라 디바이스 기준 — 블로그 사례에 따르면 TPU v7 칩 하나는 TensorCore 두 개를 PyTorch에 디바이스 두 개로 노출한다. v7x에서 --nproc_per_node를 칩 수로 적으면 절반만 쓰게 된다.
• 랭크별 분기 코드는 느려진다 — 기본 설정(Skinny Collectives)에서는 collective 연산을 각각 별도 모듈로 컴파일해 랭크 간 불일치를 막는 대신 계산·통신 겹치기를 못 해 종단 성능이 30~40% 떨어질 수 있다. get_rank() 조건 분기가 없는 구간만 @spmd_safe로 표시해 이 비용을 줄인다. 분기가 있는 구간에 붙이면 교착이나 메모리 손상이 생길 수 있다.
• FSDP 종류 확인 — 기능 현황 문서에서 SimpleFSDP는 검증 완료, FSDPv2는 실험 단계(성능 낮음)로 분류돼 있다.
• FlexAttention — 블로그의 TorchTitan 이식 사례는 아직 지원되지 않는 FlexAttention을 F.scaled_dot_product_attention으로 바꿨다.
• RL 학습·추론 같은 디바이스 공유 — TPU 디바이스는 사용자 공간 프로세스 하나만 접근할 수 있어 학습과 롤아웃을 한 디바이스에 함께 올리는 구성은 기본으로 지원되지 않는다.

Hugging Face 모델과 서빙 스택에서 이전 결과 확인하기

Hugging Face 블로그에 따르면 Transformers·Diffusers는 TorchTPU 위에서 Optimum TPU 없이 동작한다. from_pretrained에 device_map="tpu"를 넘기고 입력에 .to("tpu")를 붙이면 된다. 다만 같은 글에서 TPU로 돌린 Transformers 테스트의 통과율은 96.5%로, 모든 모델이 그대로 동작한다고 보장하지는 않는다.

from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "Qwen/Qwen3.5-9B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="tpu")

inputs = tokenizer("TPU 이전 점검용 문장", return_tensors="pt").to("tpu")
outputs = model.generate(**inputs, max_new_tokens=32)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

서빙은 vLLM과 SGLang이 TorchTPU 백엔드를 제공한다. SGLang-TorchTPU 발표는 CI와 테스트 하드웨어로 Cloud TPU v7x만 명시하고, 구조화 출력과 함수 호출은 실험 기능으로 분류한다. 서빙 API에서 JSON 출력이나 도구 호출에 의존하고 있다면 이 두 기능을 이전 검증 항목에 넣어야 한다.

11월 15일 아카이브 전에 끝내 둘 이전 작업 순서

1. 사용처 목록 만들기 — 저장소에서 torch_xla 의존 코드를 모두 찾는다: grep -rnE "torch_xla|xm\.|'xla'|\"xla\"" --include=*.py .
2. 하드웨어 세대 확인 — v5e·v5p·v6e·v7x가 아니라면 세대 전환 일정부터 잡는다.
3. 새 환경에서 단일 디바이스 실행 — 대응표대로 코드를 바꾸고 Strict 모드로 상대 MSE나 평가 점수를 기존 결과와 비교한다.
4. 속도 단계 올리기 — DeferAndFuse, 버킷 패딩을 붙인 torch.compile 순으로 켜고 스텝 시간을 기록한다.
5. 분산 실행으로 확장 — 환경 변수 설정, 디바이스 수 기준 --nproc_per_node, @spmd_safe 적용 범위를 점검한다.
6. 버전 고정과 기록 — torch_tpu==0.2.0처럼 정확한 버전을 의존성 파일에 적는다. 0.x 단계라 API가 바뀔 수 있으므로 업그레이드는 평가 회귀 비교를 거친 뒤에 한다.

이전 직후 몇 주는 스텝 시간, 컴파일 횟수, 평가 점수를 함께 기록해 두면 성능 저하가 코드 때문인지 재컴파일 때문인지 구분할 수 있다. 어떤 지표를 어떤 단위로 남길지는 LLMOps 모니터링 항목과 구현의 로그 스키마를 참고한다.