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와의 실제 속도 비교, 그리고 독립적인 성능 재현

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


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

2026-10-12

엔트로피, 그대는 왜 엔트로피인가요?

파이토치로 분류 모델을 만들다 보면 손실 함수로 거의 언제나 교차 엔트로피를 고르게 된다. 그런데 이름이 묘하다. 모델의 예측이 정답과 얼마나 다른지를 잴 때 왜 물리 시간에 배운 단어가 붙은 손실 함수를 사용할까? 엔트로피의 원래 의미를 우리말로 옮기면 '엇갈린 무질서의 정도' 쯤이 되는데 도무지 어울리지 않는다.

이유를 굳이 알 필요는 없지만, 사정을 알고 나면 왜 \(-\log p\)를 손실로 쓰는지 납득되고, 학습 시작 직후의 손실이 얼마여야 정상인지를 알 수 있게 된다.

1. 열역학의 엔트로피

무질서가 아니라 경우의 수

고등학교 물리 시간에 엔트로피는 보통 '무질서한 정도'로 가르친다. 틀린 설명은 아니지만 오해하기 딱 좋다. 정확히 말하자면 '같은 겉모양을 만들어낼 수 있는 내부의 경우의 수는 몇 가지인가'다.

책상 위에 동전 100개가 있다고 하자. 전부 앞면인 상태를 만드는 방법은 한 가지뿐이다. 반면 앞면 50개, 뒷면 50개인 상태는 어느 동전이 앞면인지에 따라 경우의 수가 \(10^{29}\) 가까이 된다. 겉으로 보이는 상태는 '앞면의 개수' 하나라는 것과 비교하면 하늘과 땅 차이다.

그래서 동전 100개를 흔들어 흩뿌리면 반반에 가까운 상태가 나온다. 무질서를 추종하려는 힘이 따로 있어서가 아니라 그렇게 될 경우의 수가 압도적으로 많기 때문이다.

왜 로그인가?

이 경우의 수를 엔트로피로 정의한 사람이 루트비히 볼츠만이다.

\[ S = k \log W \]

여기서 \(W\)가 경우의 수고, \(k\)는 상수다. 그 사이에 있는 로그에 주목해 보자. 두 사건을 나란히 놓고 보면 전체 경우의 수는 더해지지 않고 곱해진다. 엔트로피를 물리량으로 정의해 사용하려면 양이 두 배가 될 때 엔트로피도 두 배가 되면 계산이 편해진다. 로그의 역할이 바로 이것이다.

\[ \log(W_1 \times W_2) = \log W_1 + \log W_2 \]

로그를 사용하면 곱을 합으로 바꿀 수 있다. 이 원리를 기억해 두자. 뒤에서 다시 등장한다.

2. 클로드 섀넌의 엔트로피

뻔한 소식과 뜻밖의 소식

1948년 클로드 섀넌은 통신을 연구하다 이런 질문을 던졌다.

이 메시지는 얼마나 많은 정보를 담고 있을까?

시작은 단순하다. 뻔한 소식은 정보가 적고 뜻밖의 소식은 정보가 많다. '내일 해가 뜬다'는 말 속에는 새로운 정보라고 할 만한 내용이 들어 있지 않다. 만약 10월에 '내일 눈이 온다'는 말을 듣는다면 매우 놀라울 것이며, 그만큼 많은 정보가 담겨 있다.

그래서 섀넌은 어떤 사건의 정보량을 그 사건이 일어날 확률의 역수로 정의했다.그런데 여기에도 로그를 사용했다.

\[ \text{정보량} = \log \frac{1}{p} = -\log p \]

이유는 볼츠만과 같다. 독립인 두 사건이 함께 일어날 확률은 곱해지지만, 정보량 역시 더하면 여러모로 편리하기 때문이다. 예를 들어 동전을 한 번 던져 얻는 정보가 1비트일 때 두 번 던져 얻는 정보는 2비트가 되면 받아 들이기도, 계산하기도 편하다. 확률은 \(\frac{1}{2}\)과 \(\frac{1}{2}\)를 곱한 \(\frac{1}{4}\)가 되지만, 로그를 씌우면 \(\log \frac{1}{2}\)와 \(\log \frac{1}{2}\)를 더한 \(2 \log \frac{1}{2}\)가 된다.

평균적으로 얼마나 놀랐을까?

이제 사건 하나 대신 분포 전체를 생각해보자. 분포 전체를 놓고 평균 정보량을 구한 것이 바로 정보 엔트로피다. 다음과 같이 정의된다.

\[ H(p) = -\sum_{i} p_i \log p_i \]

각 사건의 정보량 \(-\log p_i\)에 그 사건이 일어날 확률 \(p_i\)를 곱해 모두 더한 값이다. 말하자면 놀라움의 평균이다. 앞면만 나오는 동전은 결과가 뻔하니 엔트로피가 0이고, 평범한 동전은 가장 예측하기 어려우니 엔트로피가 최대가 된다.

3. 두 엔트로피가 만나는 지점

이제 볼츠만과 셰년의 엔트로피를 나란히 놓아 보자.

\[ S = k \log W \]

\[ H = -\sum_{i} p_i \log p_i \]

이때 경우의 수가 \(W\)가지이고 모두 똑같이 나올 법하다면 각각의 확률은 \(\frac{1}{W}\)다. 섀넌의 엔트로피에 그대로 넣어 보면 이렇게 된다.

\[ H = -\sum_{i=1}^{W} \frac{1}{W} \log \frac{1}{W} = \log W \]

볼츠만의 식에서 상수 \(k\)만 뺀 모양인데, 우연이 아니다. 둘 다 '가능한 경우가 얼마나 많은가'를 로그로 계산한다. 물리학은 그것을 분자의 배치에 적용했고 정보 공학은 메시지에 놀랄 가능성에 적용했을 뿐이다.

폰 노이만의 오지랖에 관한 풍문

섀넌이 여기에 어떤 이름을 붙일지 고민하다 존 폰 노이만에게 물었더니 이렇게 답했다고 한다.

엔트로피라고 부르게. 첫째로 자네 함수는 통계역학에서 이미 그 이름으로 쓰이고 있고, 둘째로 엔트로피가 정확히 무엇인지는 아무도 모르니 논쟁에서 늘 유리할 걸세.

널리 알려진 이야기지만 섀넌 본인의 기록이 아니다. 사실 여부는 불명확하지만, 사실이 아니라 해도 식이 이름에 걸맞았기에 나온 이야기였을 것이다.

4. 그래서 교차 엔트로피는

두 개의 분포를 가진다.

지금까지는 분포 하나를 가진 엔트로피를 이야기했다. 그런데 분류 모델의 경우 분포가 둘이다. 하나는 정답 분포 \(p\)로, 정답이 1번 클래스라면 \([0, 1, 0, \dots]\)이다. 다른 하나는 모델이 소프트맥스 함수를 통해 계산한 예측 분포 \(q\)다.

교차 엔트로피는 정답 분포를 기준으로 삼아 모델의 예측이 얼마나 놀라운 결과인지를 계산한다.

\[ H(p, q) = -\sum_{i} p_i \log q_i \]

엔트로피 식에서 로그 안쪽만 \(q\)로 바뀌었다. 이름이 '교차'인 이유가 이것 때문이다. 한쪽 분포로 가중치를 매기고 다른 쪽 분포로 놀라움 정도를 계산하는 식으로 교차한다.

정답이 원-핫 벡터라면 수식은 매우 단촐해진다. 정답 클래스 이외에는 \(p_i\)가 모두 0이기 때문이다.

\[ H(p, q) = -\log q_{\text{정답}} \]

이것이 \(-\log p\)를 손실로 사용하는 이론적 설명이다. 해설하자면 모델이 정답 클래스에 높은 확률을 준 경우 별로 놀랍지 않고(즉, 오차가 작고), 낮은 확률을 준 경우 놀랍다(즉, 오차가 크다)는 의미다.

엔트로피를 둘로 나누면

교차 엔트로피는 다시 두 조각으로 나뉜다.

\[ H(p, q) = H(p) + D_{\mathrm{KL}}(p \,\|\, q) \]

앞쪽 \(H(p)\)는 정답 자체의 엔트로피다. 데이터가 정해지는 순간 결정되므로 모델과 무관한 고정값이다. 뒤쪽 \(D_{\mathrm{KL}}(p \,\|\, q)\)가 두 분포가 얼마나 다른지를 나타낸다. 코드로 확인해 보자.

import torch

p = torch.tensor([0.7, 0.2, 0.1])      # 정답 분포
q = torch.tensor([0.5, 0.3, 0.2])      # 예측 분포

H  = -(p * p.log()).sum()              # 0.801819
KL = (p * (p.log() - q.log())).sum()   # 0.085123
CE = -(p * q.log()).sum()              # 0.886941

print(torch.allclose(H + KL, CE))      # True

학습으로 줄일 수 있는 것은 KL 발산쪽뿐이다. 그래서 "교차 엔트로피를 최소화한다"와 "예측 분포를 정답 분포에 가깝게 만든다"가 같은 말이 된다. 정답이 원-핫이면 \(H(p)\)가 0이라 교차 엔트로피가 곧 KL 발산이 되는데, 분류 모델에서 둘을 구분하지 않고 말하는 경우가 많은 이유가 이것이다.

5. 손실의 시작값

여기까지 오면 알 수 있는 것

갓 초기화한 모델은 아무것도 모르는 모델이다. 모든 클래스에 똑같은 확률을 준다는 뜻이다. 클래스가 \(k\)개면 각각 \(\frac{1}{k}\)이므로 손실은 \(-\log \frac{1}{k}\), 즉 \(\log k\)가 되어야 한다.

클래스 수 시작 손실 해당하는 데이터셋
2 0.693 이진 분류
10 2.303 MNIST, CIFAR-10, Fashion-MNIST
100 4.605 CIFAR-100

실제로 그런지 코드로 확인할 수 있다.

import torch
import torch.nn as nn

logits = torch.zeros(4096, 10)             # 아무것도 모르는 상태
y = torch.randint(0, 10, (4096,))

print(nn.CrossEntropyLoss()(logits, y))    # tensor(2.3026)

이 숫자가 꽤 쓸 만한 진단 도구가 된다. 시작 손실이 2.3 근처라면 정상이다. 7이나 20처럼 훨씬 크다면 초기 가중치가 지나치게 크거나 로짓에 이미 소프트맥스를 먹였을 가능성이 있다. 반대로 훨씬 작다면 데이터가 새고 있거나 정답 라벨이 잘못 들어갔을 수 있다. 기울기를 들여다보기 전에 첫 숫자 하나로 거를 수 있는 문제가 생각보다 많다.

다만 학습 로그의 첫 줄과 혼동하지 말자. 대개의 학습 루프는 1에포크를 다 돌고 나서 평균 손실을 찍는다. 그 값은 이미 수천 번 갱신된 뒤라 \(\log k\)보다 한참 낮게 나온다. 위 숫자를 보려면 학습 루프에 들어가기 전에 손실을 한 번 계산해 봐야 한다.

밑이 2가 아니다.

마지막으로 자주 헷갈리는 것 하나. 정보 이론에서는 로그의 밑을 2로 두고 단위를 비트라고 부른다. 반면 파이토치를 비롯한 딥러닝 라이브러리는 자연로그를 쓰고 내트(nat)라는 단위를 사용한다.

그래서 클래스가 10개일 때 시작 손실은 \(\log_2 10 = 3.32\)가 아니라 \(\ln 10 = 2.30\) 이다. 물론 어느 쪽이든 이 차이는 상수곱으로 해결되므로 최적화 결과는 달라지지 않는다.

6. 이름값을 한다

엔트로피라는 이름은 물리학에서 빌려 온 것이지만 억지로 가져다 붙은 비유가 아니다. 경우의 수를 로그로 계산한다는 기본 접근 방식이 같았고, 곱을 합으로 바꾸기 위해 로그를 사용한다는 로그의 용도도 같다. 폰 노이만의 일화가 사실이든 아니든 이름값은 하는 셈이다.

2026-09-29

bf16과 MPS에서의 bf16 지원(2)

2. 애플 실리콘에서는 왜 쓸 수 없었는가?

CUDA가 없는 하드웨어

한때 맥에 NVIDIA 그래픽 카드가 들어가던 시절이 있었다. 하지만 2014년 맥북 프로가 NVIDIA GPU를 탑재한 마지막 맥이 되었고, macOS 10.14 모하비에서 드라이버 지원이 끊기면서 남아 있던 길도 닫혔다. NVIDIA 역시 2019년 CUDA 10.2를 마지막으로 macOS 지원을 접었다.

맥에서 CUDA를 사용할 수 없다는 말은, CUDA로 하던 일을 다른 것으로 해야 한다는 뜻이다. 애플 실리콘 맥에서 그 역할을 맡은 것이 MPS(Metal Performance Shaders) 백엔드다. 

문제는 CUDA로 할 수 있는 일과 MPS로 할 수 있는 일이 같지 않다는 데 있다.

어디에 무엇이 있는가

[ Application / Framework Layer ]
                    PyTorch User Code (Python)
                               │
                               ▼
[ PyTorch C++ Core (LibTorch / ATen) ]
             Dispatcher (CPU / CUDA / MPS 연산 분기)
                               │
                               ▼
┌-------------------------------------------------------------┐
│ [ PyTorch Internals: MPS Backend ]                          │
│  - aten/src/ATen/mps        (장치, 스트림, 메모리 할당기)   │
│  - aten/src/ATen/native/mps (연산 구현, 커스텀 Metal 커널)  │
│  - torch/csrc/mps           (torch.mps 파이썬 바인딩)       │
└-----------------------------┬------------------------------┘
                               │ (호출)
═════════════════════════════════════════════════════════════════
[ macOS System Frameworks ]    │
                               ▼
            ┌---------------------------------------┐
            │ MPSGraph (연산 그래프 융합 및 최적화) │
            │                 │                    │
            │                 ▼                    │
            │ MPS (Metal Performance Shaders)       │
            └-----------------┬--------------------┘
                               │
                               ▼
[ Low-level Graphics & Compute Layer ]
           Metal Framework / Metal Shading Language (MSL)
                          Metal Driver
                               │
                               ▼
[ Hardware Layer ]
                     Apple Silicon GPU Core
  

파이토치에서 사용하는 MPS 백엔드는 애플이 제공하는 프레임워크 위에 얹혀 있다. CUDA처럼 하드웨어 바로 위에서 커널을 작성하는 구조가 아니다.

제약은 어디에 있는가

이 구조를 이렇게 정리할 수도 있다. "애플이 제공하지 않는 기능은 파이토치도 제공할 수 없다." 절반만 맞는 말이다.

파이토치는 MPSGraph에만 의존하지 않는다. 자체 Metal 셰이더를 직접 작성해 애플의 고수준 프레임워크를 우회하기도 하며, 최근에는 그 비중을 늘리고 있다. 올 7월의 2.13 릴리스 노트에는 bernoulli, native_dropout, uniform 같은 연산을 MPSGraph에서 네이티브 Metal 커널로 옮겼다는 항목이 줄줄이 들어 있다. 애플이 만들어 준 그래프 엔진을 쓰는 대신 직접 셰이더를 쓰겠다는 선택이다.

그래서 진짜 제약은 한 단계 아래에 있다. 셰이딩 언어와 하드웨어가 지원하지 않는 것은 어느 경로로도 우회할 수 없다. 그래프 엔진에 없으면 셰이더를 직접 쓰면 되지만, 셰이딩 언어에 자료형이 없으면 셰이더를 쓸 방법 자체가 없다.

무엇이 없었고, 언제 생겼는가

bf16이 정확히 그런 경우였다.

애플 실리콘에서 bf16을 쓸 수 없었던 것은 파이토치의 게으름이나 버그 때문이 아니었다. MPS 백엔드 아래의 프레임워크와 저수준 계층에 bfloat 자료형이 아예 없었기 때문이다. 애플이 제공하는 API에는 fp16에 해당하는 half와 fp32에 해당하는 float가 있었지만, bf16의 설계에 대응하는 자료형은 없었다. MPSGraph에도 없고, 그 아래 셰이딩 언어에도 없었다. 우회할 곳이 없다는 뜻이다.

애플은 그들이 의도한 맥의 활용 방향에 따라 bfloat의 필요성을 간과했다. 애플 실리콘의 GPU와 뉴럴 엔진은 오랫동안 온디바이스 추론을 목표로 발전해 왔다. 사진 분류, 얼굴 인식, 음성 처리 같은 작업이다. 이런 용도에서는 fp16의 좁은 표현 범위가 거의 문제가 되지 않는다. 학습이 아니라 이미 학습된 모델을 실행하는 일이기 때문이다. 기울기 소실을 걱정할 필요가 없는 상황에서는 오히려 가수부가 3비트 더 넉넉한 fp16이 정밀도 면에서 유리하다.

그래서 애플은 fp16을 하드웨어와 셰이딩 언어 양쪽에서 먼저 지원해 왔다. 맥에서 큰 모델을 학습하거나 미세 조정하려는 수요는 뒤늦게 인정한 셈이다.

결국 애플은 macOS 14 소노마에서 bf16을 도입했다. Metal Shading Language 3.1에 bfloat 타입이 생기고, MPSGraph에 MPSDataTypeBFloat16이 추가되었다. 파이토치는 2.3 릴리스에서 이를 받았다. 릴리스 노트에 한 줄로 적혀 있다.

Enable bfloat16 support on MacOS 14 (#119641)

즉 애플 실리콘 맥에서 macOS 14 이상, 파이토치 2.3 이상이면 이제 bf16을 사용할 수 있다. 2024년의 일이다.

3. 그래도 아직 부족하다

문제가 bf16에만 있었던 것은 아니다. MPS 백엔드가 제공하지 않는 기능, CUDA에서만 되는 기능은 지금도 남아 있다.

목록에는 있고 구현되지 않은 연산들

파이토치는 MPS 백엔드에 구현되기를 바라는 연산을 투표로 모으고 있다(링크). 2025년 5월에 만들어진 목록이다. 이후 구현된 항목에는 취소선이 그어졌는데, 아직 다수가 남아 있다. 

남은 것을 보면 linalg_eig, linalg_qr, _linalg_eigh, linalg_matrix_exp처럼 선형대수 분해가 압도적이다. 고윳값 분해와 QR 분해, 촐레스키가 MPS의 가장 큰 공백이다.

목록에조차 없는 것

이 목록이 전부도 아니다.

<멀티모달까지 직접 구현하는 딥러닝 with 파이토치> 10-1절의 예제는 MPS 백엔드에서 다음과 같은 예외가 발생한다.

NotImplementedError: The operator 'aten::_nested_tensor_from_mask_left_aligned' is not currently implemented for the MPS device.

이 연산(_nested_tensor_from_mask_left_aligned)은 투표 집계 목록에 없다. 목록에 오르려면 셋 이상의 요청이 필요한데, 두 건의 요청만 받았기 때문이다. 그 두 표 중 하나를 남긴 사람이 <Hands-On Machine Learning>의 저자 오렐리앙 제롱이다. 용례도 똑같이 "Training a Transformer"라고 적혀 있다. 2025년 3월의 일이고, 1년 반이 지난 지금도 구현되지 않았다.

이 연산 자체는 대단한 것이 아니다. 패딩 마스크가 좌측 정렬인지 판정하는 연산일 뿐이다. 문제는 nn.TransformerEncoder의 기본 설정이 이 연산을 사용한다는 데 있다. 패딩을 건너뛰어 연산량을 줄이는 경로인데, 이를 위해 MPS에 없는 연산을 사용한다. 

다행히 우회할 방법이 있다. torch.backends.mha.set_fastpath_enabled(False)를 호출해 전역 스위치를 끄면 된다. MPS 백엔드 대신 CPU를 사용하는 것도 한 가지 쉬운 대안이 될 수 있다.

늘 한 발 늦는다

이런 사례들을 보면 MPS 백엔드 지원의 패턴을 가늠할 수 있다. 언젠가 지원되지만 늘 한 발 늦는다는 점이다. 예를 들어 지난 7월 2.13 릴리즈에 포함된 MPS 백엔드의 FlexAttention은 CUDA 백엔드는 2.5 릴리즈(2024년 10월)에 포함되었다. 

어쩔 수 없다. 현재 딥러닝 개발의 표준은 CUDA라는 사실은 부인하기 어렵다. 새 기능은 CUDA에서 나오고, MPS는 여러 릴리스 뒤에 따라온다. 개발자 친화적인 하드웨어와 운영체제라 해도, 파이토치로 딥러닝을 하는 사람에게 맥을 선뜻 권하기 어려운 이유가 이것이다.

애플이 직접 내놓은 대안으로 MLX가 있지만 파이토치만큼 널리 쓰이지는 않는다.

기대감

그래도 조금씩 좋아지고 있다.

올 7월의 2.13 릴리스에는 MPS 관련 항목이 스무 개 넘게 들어갔다. CUDA에 2.5부터 있던 FlexAttention이 MPS에서도 동작하게 되었고, 릴리스 노트는 "희소 패턴에서 SDPA 대비 최대 12배"라고 적었다. 늦게 왔지만 올 때는 제대로 온 셈이다.

더 눈에 띄는 것은 앞서 이야기한 흐름이다. 파이토치가 MPSGraph에서 벗어나 자체 Metal 커널을 늘려 가고 있다. 애플이 만들어 주기를 기다리는 대신 직접 만드는 쪽으로 무게가 옮겨 가는 중이다. 그렇게 되면 "애플이 제공하지 않으면 못 한다"는 제약이 닿는 범위도 그만큼 줄어든다.

언젠가는 애플 실리콘의 좋은 하드웨어를 파이토치에서 온전히 활용할 수 있는 시기가 올 것이라는 기대를 가져 본다.


2026-09-26

bf16과 MPS에서의 bf16 지원(1)

리눅스에서 잘 동작하는 파이토치 코드를 맥에서 실행하면 의도대로 동작하지 않는 경우가 종종 있다. 특히 연산 장치를 사용하는 경우에 자주 발생하는데, 이는 NVIDIA의 CUDA와 Apple Silicon의 MPS의 지원 수준이 다르기 때문이다.

이 지원 수준은 파이토치가 얼마만큼 지원하느냐도 중요하지만, 파이토치 같은 하드웨어의 성능을 최대한 활용하는 프레임워크의 경우 네이티브 API가 무엇을 지원하느냐도 중요하다. 그 중 하나의 이야기다.

1. bfloat16은 무엇일까?

16비트 실수를 표현하는 두 가지 방법

부동소수점 수는  부호부(1비트), 지수부, 가수부로 메모리를 비트별로 나눠 사용한다. 여기서 지수부는 수의 크기(자릿수)를 정하고, 가수부는 정밀도(유효숫자)를 정한다. 파이토치의 기본 실수 자료형은 torch.float32(fp32)인데, fp32의 구조는 다음과 같다.
fp32: SEEEEEEEEMMMMMMMMMMMMMMMMMMMMMMM   부호(S)*1 + 지수(E)*8 + 가수(M)*23

메모리를 절반만 사용하는 16비트 실수 자료형에는 torch.float16(fp16)뿐 아니라 torch.bfloat16(bf16)도 있다.  두 자료형의 내부 구조는 다음과 같다.

fp16: SEEEEEMMMMMMMMMM    부호(S)*1 + 지수(E)*5 + 가수(M)*10

bf16: SEEEEEEEEMMMMMMM    부호(S)*1 + 지수(E)*8 + 가수(M)*7

fp16은 지수부를 5비트로 줄이고 가수부를 최대한 늘려 정밀도를 지키고 범위를 줄인 반면,  bf16은 반대의 선택을 하고 있다. 여기서 재미있는 것은 bf16은 지수부가 8비트로 fp32와 같고, 마지막 가수부만 23비트에서 7비트로 줄었다는 점이다. 그래서 변환 과정은 뒤의 16비트만 잘라내면 되어 비용이 매우 낮다.

실제 값이 어떻게 될까?

파이토치로 직접 확인해 보자.
import torch

for name, dt in [('fp32', torch.float32),
                 ('fp16', torch.float16),
                 ('bf16', torch.bfloat16)]:
    i = torch.finfo(dt)
    print(f'{name}: max={i.max:.6g}  min={i.tiny:.6g}  eps={i.eps:.6g}')
fp32: max=3.40282e+38  min=1.17549e-38  eps=1.19209e-07
fp16: max=65504        min=6.10352e-05  eps=0.000976562
bf16: max=3.38953e+38  min=1.17549e-38  eps=0.0078125

결과를 확인해 보면 bf16의 범위는 fp32와 같지만 정밀도는 훨씬 떨어진다. 반대로 fp16은 범위와 정밀도 모두를 적당히 깎아서 타협한 결과물이다. 하지만 이 타협의 결과는 썩 좋은 편이 아니다. 범위의 문제가 발목을 잡는다.

오버플로(위로 넘칠 때)와 언더플로(아래로 넘칠 때) 모두 문제가 되지만, 둘 중 언더플로의 문제가 더 크다. 오버플로는 그나마 좀 낫다. fp16은 값이 65504를 넘는 순간 inf가 되고, 이를 사용하는 연산의 결과는 nan이 된다. 반면 bf16은 fp32의 범위를 거의 그대로 다 받는다. 그런데 나은 것이 이정도다. 언더플로는 inf고, nan이고 그런 것 없이 그냥 값이 0이 되어 버린다. 아무런 경고도, 예외도 없이 말이다. 무섭지 아니한가?

그래서 하는 짓: 손실 스케일링

학습 중에 계산되는 기울기는 fp16의 언더플로 영역을 자주 드나든다. 이때 값이 완전히 사라지고 0이 되면, 해당 파라미터는 더 이상 변하지 않는다. 학습 곡선이 어찌되었건 모델의 여기저기에서 학습이 멈추는 일이 벌어진다.

이를 피하기 위해서 '손실 스케일링'이라는 것을 한다. 역전파 전에 손실에 큰 수를 곱해서 기울기를 fp16 범위 안으로 끌어 올리고, 최적화 전에 다시 나누는 방식이다. GradScaler가 하는 일이 바로 이것인데, 오버플로가 발생하면 반대 방향의 일도 함께 처리한다.

한편 bf16은

bf16은 범위가 fp32와 거의 같으므로 애초에 오버플로와 언더플로를 걱정할 일이 크게 줄어든다. 반면 정밀도가 낮은 것은 생각만큼 큰 문제가 되지 않는다. '확률적' 경사 하강법 자체가 확률로 인한 노이즈가 끼어들기 때문이다.

fp16과 bf16의 차이를 정리하면 다음과 같다.

항목fp16bf16
범위좁음 (최대 65504)fp32와 같음
정밀도상대적으로 높음낮음 (유효숫자 2~3자리)
학습손실 스케일링 필요대부분 그냥 가능
추론유리함무난함
fp32 변환비트 재배치 필요뒤를 잘라내면 됨

bf16의 유래

이름에 붙은 b는 Brain을 뜻한다. 구글 브레인이 TPU 전용 자료형으로 설계한 형식으로, 그래서 b가 붙어 있다.

학습의 대상이 되는 모델이 점점 커지면서 널리 보급되었다. 엔비지아가 암페어(A100, RTX30 시리즈)부터 하드웨어 레벨로 지원했고, 이후 큰 모델을 학습할 때 사실상 기본값이 되었다. 허깅페이스 예제 코드에도 dtype=torch.bfloat16을 쉽게 확인할 수 있다.

그런데 어떤 생태계에서는 표준이나 마찬가지인 bf16은 Apple Silicon의 MPS에서는 이야기가 달라진다.

(2편에서 계속)

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

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