Showing posts with label 파이토치. Show all posts
Showing posts with label 파이토치. Show all posts

2026-10-22

코랩 TPU에서 파이토치, 이제는 쓸 만 할까? - TorchTPU의 등장

<멀티모달까지 직접 구현하는 딥러닝 with 파이토치>의 코랩 안내 문서에는 이런 문장이 있다. "TPU는 고르지 말 것." 코랩의 런타임 유형 메뉴에서 TPU를 선택하면 GPU처럼 하드웨어 가속기의 장점을 누릴 수 있을 것처럼 보이는데, 정작 예제 노트북의 일부는 TPU 런타임에서 제대로 동작하지 않는다.

그런데 2026년 10월 20일, 구글과 메타가 TPU에서 파이토치를 네이티브로 실행하는 백엔드 TorchTPU를 공개했다. 코랩 런타임에 이미 적용되었다는 공지까지 올라왔다. 그렇다면 이제 마음 편하게 TPU 런타임을 사용해도 될까? 이 글은 TorchTPU가 무엇인지, CUDA와 무엇이 다른지, 그리고 코랩 무료 TPU에서 잠시 돌려 본 결과를 정리한 것이다.

참고로 이 글은 2026년 10월 22일 기준이다. 코랩 런타임과 TorchTPU 모두 빠르게 바뀌고 있어서 읽는 시점에 따라 상황이 다를 수 있다.

1. 코랩에서 TPU를 피하라고 한 이유

얼핏 보면 GPU보다 좋아 보인다.

TPU는 구글이 만든 딥러닝 전용 칩이다. 범용 그래픽 카드인 GPU보다 딥러닝에 특화되어 있다고 하니 당연히 더 좋을 것 같고, 런타임 유형 메뉴에서 TPU를 골라 보고 싶은 유혹도 생길 만하다.

그런데 <멀티모달까지 직접 구현하는 딥러닝 with 파이토치>의 코랩 안내 문서(저장소의 articles/prolog/google-colab.md)는 반대로 말한다. TPU는 고르지 말고, GPU를 배정받지 못했다면 차라리 CPU를 사용하라고 안내한다.

TPU를 사용하면 오류 없이 느려진다

이유는 이렇다. 책의 학습 코드는 일반적인 파이토치 방식으로 쓰여 있다. 반면 TPU는 연산을 모아 두었다가 한꺼번에 실행하는 방식이라 TPU 전용 코드가 따로 필요하다. 범용 환경을 기준으로 하는 입문서에서 다루기 어려운 내용이다. 그래서 TPU 런타임에서 책의 예제 노트북을 실행하면 오류가 나는 대신 극도로 느려지거나 멈춘 것처럼 보인다.

오류가 났다면 차라리 나았을 것이다. 오류 메시지라도 있으면 무엇이 문제인지 짚어 볼 수 있으니까. 아무 말 없이 느려지거나 멈추는 상황이 가장 고약하다.

여기서 말한 "연산을 모아 두었다가 한꺼번에 실행하는 방식"은 2020년에 나온 PyTorch/XLA의 Lazy Tensor 방식이다. 그리고 이번에 발표된 TorchTPU의 설계는 이 방식에서 완전히 벗어나 있다.

2. TorchTPU가 나타났다

TorchTPU는 무엇일까?

이름 때문에 생길 수 있는 오해부터 바로잡자. TorchTPU는 새로운 TPU 칩이 아니다. TPU는 구글이 오래전부터 만들어 온 하드웨어이고, TorchTPU는 그 위에서 파이토치를 돌리기 위한 소프트웨어(백엔드)다. 파이토치에서 device="cuda"가 엔비디아 GPU를 가리키듯, device="tpu"가 TPU를 가리키게 해 주는 부분이다.

PyTorch 모델 코드          model.to("tpu")
   └── TorchTPU 백엔드      ← 이번에 발표된 것 (소프트웨어)
          └── XLA 컴파일러 (StableHLO, Pallas)
                 └── TPU 칩 (v5e, v6e, v7 …)   ← 하드웨어, 기존 그대로

TorchTPU는 구글과 메타가 2025년 중반부터 함께 만든 오픈소스 프로젝트다. 2026년 10월 20일 PyTorch Conference에서 공개되었고 소스는 google-pytorch/torch_tpu(Apache-2.0)에 있다. 구글 클라우드 문서는 TorchTPU가 기존의 PyTorch/XLA를 대체한다고 밝히고, 새 워크로드는 TorchTPU로 시작하라고 권한다.

무엇이 달라졌는가

핵심은 다섯 가지다.

  • 이식이 쉽다. model.to("cuda")를 model.to("tpu")로, 분산 백엔드 nccl을 tpu_dist로 바꾸는 수준이다. 다만 PyTorch 블로그도 이것만으로 TPU의 최고 성능이 나오지는 않는다고 인정한다.
  • eager 모드를 최우선으로 지원한다. eager 모드는 파이토치가 코드를 한 줄씩 읽는 대로 바로 실행하는 방식이다. 우리가 평소 쓰는 파이토치의 기본 동작이라서 중간 결과를 print로 찍어 보고 오류가 난 줄을 바로 찾아갈 수 있다. 앞에서 언급한 연산을 모아 두었다가 한꺼번에 실행하는 방식(lazy)과 반대다. 예전 PyTorch/XLA는 이 방식을 지원하지 않았고, 그래서 책의 예제 노트북이 오류 없이 느려지는 일이 벌어졌다. TorchTPU는 eager를 세 가지로 제공한다.
    • Strict(기본): 연산을 하나씩 TPU에 보낸다. GPU에서 쓰는 방식과 같다.
    • Debug: 연산마다 TPU의 실행이 끝나기를 기다린다. 느리지만 오류 위치를 정확히 찾을 수 있다.
    • DeferAndFuse: 겉보기에는 eager지만 내부에서 연산을 모았다가 묶어서 컴파일한다. Strict보다 50% 이상 빠를 수 있다고 한다.
  • torch.compile과 통합된다. torch.compile(model, backend="tpu")로 쓰며, Lazy Tensor가 아니라 TorchDynamo 기반이다. TorchDynamo는 코드에서 파이토치 연산을 뽑아 그래프로 묶어 주는 torch.compile의 앞단이고, TorchTPU는 그 그래프를 받아 TPU용으로 컴파일하는 백엔드 역할을 한다.
  • 런타임에 JAX 의존이 없다. 내부적으로는 XLA 컴파일러 스택(StableHLO)을 쓰지만 사용자가 JAX를 알 필요는 없다.
  • 생태계와 연결된다. Hugging Face Transformers v5, Diffusers, vLLM, SGLang, TorchTitan, Ray와의 연동이 검증되었다고 한다. 상위 5,000개 Hugging Face 모델을 대상으로 한 테스트에서 4,917개 모델 중 98%가 크래시 없이 통과했고, 대표 구조 657개 중 97%가 GPU와 수치 일치 기준을 통과했다. 모두 구글과 메타가 직접 밝힌 수치다.

3. CUDA와 비교하면

속도는 비교하기 이르다

먼저 짚어 둘 것이 있다. 공식 문서에는 TorchTPU와 CUDA GPU의 직접적인 속도 비교가 없다. "TPU가 더 빠르더라"는 말을 하고 싶지만 근거로 삼을 자료가 없다. 대신 설계가 어떻게 다른지는 비교할 수 있다.

설계의 차이

항목 PyTorch on CUDA PyTorch on TPU (TorchTPU)
기본 연산 정밀도 FP32(highest) BF16 입력 + FP32 누적(medium)
행렬곱 가속 Tensor Core (TF32) 시스톨릭 어레이 MXU (bfloat16)
정밀도를 높일 때 비용 약 1~2배 high 약 3배, highest 약 6배
집합 통신 NCCL StableHLO 집합 연산(XLA)
동시 실행 여러 SM이 커널을 병렬 실행 TensorCore는 단일 스레드, DMA와 SparseCore로 통신·전송을 겹침
입력 크기 변화 자유로움 정적 shape 전제. torch.compile은 현재 dynamic=False 고정, bounded dynamism은 개발 중
프로세스 모델 디바이스당 1프로세스 디바이스당 1프로세스. 한 TPU는 한 프로세스만 접근 가능(RL colocation 불가)

첫 줄이 눈에 띈다. CUDA에서는 기본이 순수 FP32인데 TorchTPU는 기본이 BF16이다. 앞선 글 bf16과 MPS에서의 bf16 지원(1)에서 bf16의 b가 Brain이고 TPU를 위해 만든 형식이라고 했던 것을 떠올려 보자. 그 TPU가 기본값으로 이 형식을 쓰는 것이니 이름값을 하는 셈이다.

발표된 성능 수치

PyTorch 블로그 원문에는 숫자가 잔뜩 나온다. 모두 구글과 메타가 내놓은 것이다. 숫자를 보기 전에 용어 두 개만 짚고 가자.

  • MFU: 칩이 이론상 낼 수 있는 최대 계산 능력 중 모델 계산에 실제로 쓴 비율이다. 공장에 비유하면 "설비가 하루에 만들 수 있는 최대 생산량 대비 실제 생산량"이다. 100%는 현실적으로 불가능하고, 대규모 학습에서는 30~50%대면 준수하다고 보는 경우가 많다.
  • TPS(tokens per second): 초당 처리하는 토큰 수다. 토큰은 모델이 글을 읽는 단위(단어 조각)이므로 클수록 빠르다.

이제 숫자를 하나씩 풀어 보자.

발표된 수치 쉽게 말하면
1,024칩에서 8,192칩까지 늘려도 계산 코어(TensorCore, 칩 하나에 2개) 하나당 처리량이 2,236에서 2,104 TPS로 약 6% 줄어드는 데 그쳤다(Qwen3 MoE 235B 학습, MFU 27~29%) 일꾼을 8배로 늘렸는데 1인당 생산량은 94% 정도를 유지했다는 뜻이다. 표의 숫자로 전체 처리량을 계산하면 약 458만 TPS에서 약 3,447만 TPS로 약 7.5배가 된다. 칩을 늘린 만큼 거의 그대로 빨라진다는 말이다
Llama 3 405B를 1,024칩으로 학습했을 때 칩당 초당 500토큰, MFU 50% 이상 4,050억 개의 파라미터를 가진 거대 모델에서 칩 성능의 절반 이상을 계산에 썼다. 다만 앞의 MFU 27~29%와는 모델과 설정이 달라 직접 비교할 수 없다
같은 64칩에서 TorchTitan+TorchTPU가 1,828 TPS, JAX/MaxText가 1,937 TPS 같은 TPU에서 파이토치로 짠 코드가 구글의 기존 도구인 JAX의 약 94% 속도를 낸다. 프레임워크가 바뀌어도 속도 손실이 크지 않다는 뜻이라 이 글의 맥락에서 가장 의미 있는 숫자다
외부 벤치마크 SemiAnalysis InferenceX에서 vLLM-TorchTPU가 100만 토큰당 2.01센트 모델이 만든 글 100만 토큰의 비용이 약 2센트, 1달러로 약 5천만 토큰이다. 원문에 따르면 9월의 3.6센트에서 약 44% 내려갔다. 참고로 이것은 학습이 아니라 추론(서비스) 비용이다

앞의 세 항목은 모델을 만드는 학습, 마지막 항목은 만들어진 모델을 쓰는 추론에 대한 수치라는 점을 유의해서 보자.

공개 직후에 쏟아진 기사들도 같은 점을 짚는다. Techzine은 TorchTPU가 "구글 칩이 모든 워크로드에서 엔비디아 GPU보다 빠르거나 싸다는 것을 자동으로 의미하지는 않는다"고 정리한다. 여기서 얻을 결론은 "TPU가 더 빠르다"가 아니라 "TPU로 갈아타는 장벽이 낮아졌다"다.

4. TorchTPU가 해결하는 PyTorch 문제들

네 가지 한계

PyTorch 블로그가 PyTorch/XLA의 한계로 든 네 가지에 TorchTPU는 이렇게 대응한다.

PyTorch/XLA의 한계 TorchTPU의 접근
컴파일 영역을 소스에서 명시해야 함 일반 파이토치 코드가 그대로 실행됨. torch.compile은 Dynamo 기반
컴파일 지연이 큼 3단계 컴파일 캐시(프로세스, 호스트, 원격 저장소)
eager 모드 미지원 Strict, Debug, DeferAndFuse 세 가지 eager
새 분산 기능이 뒤처짐 c10d::Backend 수준에서 집합 통신을 구현해 DDP, FSDP, DTensor 등을 그대로 지원

이 중 앞에서 언급한 "TPU 전용 코드가 필요하다"는 문제는 첫 번째 항목에 해당한다.

아직 남은 한계

물론 모든 문제가 풀린 것은 아니다.

  • FlexAttention은 아직 지원되지 않는다.
  • 집합 연산의 기본 모드(Skinny Collectives)는 종단 간 성능을 30~40% 깎을 수 있어서, @spmd_safe나 compilation handshake로 회피해야 한다.
  • 제3자의 저장소 감사에 따르면 C++로 등록되지 않은 연산이 365개 남아 있다. 목록에는 "임시"라는 표기가 붙어 있다.
  • 성능 수치는 대부분 자체 발표이고 독립적인 재현은 아직 많지 않다.

5. 코랩에서 직접 돌려 보니...

당장 사용할 수 있다

2026년 10월 20일, 코랩 공식 저장소에 공지가 올라왔다. "Colab TPU 런타임이 TorchTPU, PyTorch 2.14.1 등으로 업그레이드되었다"는 내용이다. 그러면 직접 확인해 보자.

아래 실험은 모두 무료 계정에서 v5e-1로 실행했다. 런타임 유형 메뉴에서 선택할 수 있는 TPU는 v5e-1이었고 v6e-1은 비활성이었다. 선택 가능 여부는 요금제나 그 시점의 가용 자원에 따라 달라질 수 있다.

일단 돌아간다

import sys, torch
print(sys.version); print(torch.__version__)
x = torch.randn(4, 4, device="tpu")
print((x @ x).cpu())

Python 3.13.16, PyTorch 2.14.1+cpu 환경에서 에러 없이 실행되었고 행렬곱 결과도 CPU로 돌아왔다. 설치된 torch_tpu는 0.2.0rc4, 정식 버전(0.2.0)이 아닌 릴리스 후보(rc)였다.

학습 스텝도 동작한다

간단한 2층 MLP(입력 512, 은닉 1024, 출력 10)로 SGD 학습 스텝을 돌려 보았다.

모드 첫 스텝 이후 20스텝 손실
eager 0.19초 0.02초 2.326 → 0.443
torch.compile(backend="tpu") 1.04초 0.02초 0.406 → 0.115

컴파일 모드는 eager에 이어 같은 모델을 계속 학습한 것이라 손실이 이어진다. 여기서 알 수 있는 것은 다음 세 가지다.

  • 순전파, 역전파, 옵티마이저로 이어지는 학습 루프가 eager와 compile 양쪽에서 돌고 손실이 줄었다.
  • compile은 첫 호출에 시간이 더 필요하다(1.04초). 이후 속도 차이는 이렇게 작은 모델에서는 보이지 않는다. Hugging Face의 소개글도 compile의 첫 호출이 느리다고 밝혔다.
  • 코드에서 .grad를 건드린 적이 없는데 "non-leaf Tensor의 .grad" 경고가 loss.item() 줄에서 떴다. 학습에는 영향이 없어 보이지만 원인은 확인하지 못했다.

이 실험이 보여 주지 못하는 것

작은 모델 하나를 돌렸을 뿐이다. 큰 모델, 메모리 한도, 장시간 학습, 무료 티어의 사용 시간 제한, T4 GPU와의 속도 비교는 확인하지 않았다. "돌아간다"와 "쓸 만하다" 사이에는 꽤 먼 거리가 있다.

6. 그래서 예제 노트북은 어쩌란 말이지?

앞선 안내에서는 "TPU는 전용 코드가 필요하고 오류 없이 극도로 느려지니 고르지 말라"고 했다. 이 안내가 어떻게 바뀌어야 하는지 나눠서 살펴보자.

달라진 것

  • 권고의 근거였던 "TPU 전용 코드가 필요하다"는 점은 TorchTPU가 정면으로 해결하려는 문제다. .to("tpu")로 일반 파이토치 코드가 실행된다.
  • TorchTPU가 코랩 TPU 런타임에 이미 들어와 있다.
  • 무료 계정에서 선택할 수 있는 v5e-1에서 기본 연산, 학습 스텝, torch.compile이 동작했다.

그대로인 것, 아직 모르는 것

  • 이 책의 예제 전체를 TPU 런타임에서 테스트해 보지 않았다.
  • 코랩 무료 TPU의 사용 시간 제한과 런타임 이미지의 안정성은 공식 문서에 없다. FAQ는 제공되는 TPU 종류가 "시간에 따라 바뀐다"고만 말한다.
  • 코랩에는 이전 런타임으로 돌아가는 임시 경로(fallback)가 있다. 그만큼 이 환경이 아직 바뀌는 중이라는 뜻이다.
  • 설치된 것이 정식 0.2.0이 아닌 rc 버전이었고 원인 불명의 경고도 있었다.

결론: 아직 테스트가 더 필요하다

TPU 사용을 권장하기에는 조금 이르다

TorchTPU가 "전용 코드가 필요하다"는 이유를 해결하는 방향으로 분명히 움직이고 있고, 코랩 무료 TPU에서 기본적인 학습 스텝이 도는 것도 확인했다. 하지만 이 책의 예제 전체를 테스트해 보지는 못했고, 무료 티어의 한도와 안정성도 아직 모른다. 테스트가 끝나면 그 상황을 공지하겠다. 그때까지는 TPU 사용에 주의하는 것이 좋겠다.

궁금한 독자는 실험 삼아 .to("tpu")를 시도해 볼 수 있다. 다만 어디까지나 실험으로 접근하기 바라며, 의미 있는 결과가 나왔다면 댓글 등을 통해 알려 주면 좋겠다.

앞으로 확인할 것

  • 이 책의 예제 전체를 TPU 런타임에서 실행해 본 결과
  • 무료 티어의 사용 시간 제한과 런타임 안정성
  • 정식 릴리스(0.2.0 이후)와 코랩의 공식 안내
  • T4 GPU와의 실제 속도 비교, 그리고 독립적인 성능 재현

확인이 끝나면 결과를 공지하겠다.


참고한 자료는 다음과 같다.

코랩 TPU에서 파이토치, 이제는 쓸 만 할까? - TorchTPU의 등장

<멀티모달까지 직접 구현하는 딥러닝 with 파이토치>의 코랩 안내 문서에는 이런 문장이 있다. "TPU는 고르지 말 것." 코랩의 런타임 유형 메뉴에서 TPU를 선택하면 GPU처럼 하드웨어 가속기의 장점을 누릴 수 있을 ...