TorchTPU đưa TPU vào PyTorch: chạy Hugging Face bằng `device_map="tpu"`
TorchTPU bổ sung TPU như một backend PyTorch, hỗ trợ Transformers, Diffusers và Kernels. Với TPU VM, developer có thể nạp model qua devicemap="tpu" hoặc chuyển pipeline Diffusers bằng pipe.to("tpu"), không cần chuyển code sang JAX.
Trước đây, chạy PyTorch trên Google TPU thường đồng nghĩa với torchxla hoặc chuyển sang JAX. Hugging Face giới thiệu TorchTPU để thay đổi điểm vào này: TPU được dùng như một thiết bị PyTorch, bên cạnh CPU và GPU.
Với Transformers, model có thể được nạp bằng devicemap="tpu". Với Diffusers, pipeline có thể chuyển sang TPU qua pipe.to("tpu"). Đây là lựa chọn đáng chú ý cho team đã dùng PyTorch và Hugging Face, đồng thời muốn thử hạ tầng TPU trên Google Cloud mà không phải viết lại toàn bộ code theo JAX.
TorchTPU thay đổi gì?
TPU là viết tắt của Tensor Processing Unit — bộ tăng tốc AI do Google thiết kế, tập trung vào các phép tính ma trận của mạng nơ-ron. Google Cloud cung cấp TPU từ một chip đến các cụm quy mô lớn.
Theo Hugging Face, TorchTPU bổ sung TPU như một backend PyTorch gốc, với:
PyTorch eager mode;
tích hợp torch.compile;
hỗ trợ huấn luyện phân tán;
Pallas kernels tối ưu cho TPU;
hỗ trợ trong Transformers, Diffusers và Kernels.
Trước đó, PyTorch trên TPU chủ yếu đi qua PyTorch/XLA (torchxla). Hugging Face cho biết cơ chế tensor lazy, khả năng biên dịch lại graph khi code thay đổi và việc debug đã khiến luồng này khó sử dụng hơn. Khi dùng TorchTPU với Transformers và Diffusers, Optimum TPU không còn là yêu cầu.
Nói ngắn gọn: TPU trở thành một target runtime mới cho code PyTorch hiện có, thay vì đòi hỏi một nhánh triển khai riêng bằng JAX.
Chạy Transformers trên TPU VM
TorchTPU vẫn yêu cầu phần cứng TPU hoặc TPU VM. Trên TPU VM, Hugging Face đưa ra lệnh cài đặt:
bash
pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu
Theo bài công bố, extra index này cung cấp bản CPU của torch; nếu không có, pip có thể tải bản CUDA lớn hơn dù không cần dùng CUDA trên TPU.
Ví dụ với Transformers:
python
from transformers import AutoModelForCausalLM, AutoTokenizer
modelid = "Qwen/Qwen3.5-9B"
model = AutoModelForCausalLM.frompretrained(
modelid,
devicemap="tpu"
)
tokenizer = AutoTokenizer.frompretrained(modelid)
inputs = tokenizer(
"The future of AI is",
returntensors="pt"
).to("tpu")
out = model.generate(inputs, maxnewtokens=64)
Ở đây, Transformers nạp model trực tiếp lên TPU thông qua devicemap="tpu", còn tensor đầu vào được chuyển sang "tpu" trước khi gọi generate().
Hugging Face cho biết model chạy eager mode theo mặc định. Khi cần tối ưu thêm, TorchTPU cũng tích hợp với torch.compile và hỗ trợ thiết lập phân tán.
Diffusers: chuyển pipeline bằng pipe.to("tpu")
TPU cũng là thiết bị được hỗ trợ trong Diffusers. Với eager mode, thay đổi cốt lõi là:
python
pipe.to("tpu")
Ví dụ đầy đủ từ bài viết:
python
import torch
from diffusers import StableDiffusion3Pipeline
pipe = StableDiffusion3Pipeline.frompretrained(
"stabilityai/stable-diffusion-3.5-medium",
dtype=torch.bfloat16
)
pipe.to("tpu")
pipe.transformer.compile(
backend="tpu",
fullgraph=True,
dynamic=False
)
image = pipe(
"a corgi astronaut planting a flag on a pizza moon"
).images[0]
Phần compile() là tùy chọn. Hugging Face mô tả eager mode phù hợp cho prototype hoặc debug, còn compiled mode hướng đến triển khai hiệu năng cao hơn.
Lần gọi đầu tiên sẽ thực hiện compile. Các lần gọi sau có thể tái sử dụng graph khi kích thước ảnh giữ nguyên.
Hugging Face đưa ra một số số liệu trên Cloud TPU v6e-4, với tensor parallelism 4 chip, bf16 và batch size 1. Ví dụ Qwen-Image-2.1 với prompt capybara đạt 3,5 giây ở compiled mode, so với 8 giây ở eager mode. Một ví dụ MiniMax-H3 tạo video kèm âm thanh mất 188 giây ở compiled mode và 262 giây ở eager mode. Đây là kết quả trên cấu hình và workload cụ thể của Hugging Face.
Model lớn và kernel cho MoE
Nếu model không vừa một chip, Hugging Face cho biết có thể truyền TensorParallelConfig vào frompretrained để chia transformer qua nhiều chip ngay trong lúc nạp model. Nhiều pipeline đã hỗ trợ sẵn cách làm này; những pipeline khác cần cấu hình tpplan.
TorchTPU cũng đi cùng hệ sinh thái kernel. Kernels Hub có các TPU kernel dùng lại được, viết bằng Pallas. Hugging Face cũng nhắc đến Tokamax, nơi có thể cung cấp kernel để tích hợp vào model Mixture-of-Experts.
Trong benchmark của Hugging Face trên TPU v6e-8, kernel ragged-dot cho tốc độ decode cao hơn 21,5 lần trên Qwen3-30B-A3B và 12,8 lần trên Gemma 4 26B-A4B, ở batch size 64. Con số này là so sánh giữa Hub kernel và implementation mặc định trong điều kiện thử nghiệm đã nêu, không đại diện cho mọi mô hình hay cỡ batch.
Mức độ hoàn thiện hiện tại
TorchTPU chưa phải một bức tranh hoàn toàn “xanh”. Hugging Face cho biết 96,5% số bài test chạy trên TPU đang pass. Nhóm đã báo cáo 34 issue cho TorchTPU, trong đó 19 issue đã được sửa tại thời điểm bài viết được công bố.
Việc chạy test trên TPU cũng đã phát hiện lỗi trong Transformers, bao gồm lỗi liên quan đến StaticCache và kế hoạch tensor parallel của Qwen3. Hugging Face cho biết các lỗi này đã được sửa ở upstream.
Với các team đang đánh giá TPU, điểm đáng chú ý không phải là lời hứa “đổi device là xong”, mà là khả năng bắt đầu bằng một thử nghiệm nhỏ: lấy model hoặc pipeline Hugging Face hiện có, chạy trên TPU VM, rồi đo kết quả trong điều kiện của chính workload đó.
TorchTPU hạ thấp rào cản phần mềm để thử TPU từ PyTorch. Với developer đã ở trong hệ Hugging Face, câu hỏi giờ gần hơn với: model hiện tại hoạt động thế nào khi target chuyển sang "tpu"?