Hugging Face ra mắt TorchTPU, đưa TPU thành thiết bị gốc trong PyTorch
Hugging Face công bố TorchTPU, backend PyTorch gốc cho phép lập trình viên sử dụng các Tensor Processing Unit (TPU) của Google như một thiết bị thông thường bên cạnh CPU và GPU. Bản phát hành này bổ sung hỗ trợ TPU gốc cho Hugging Face Transformers, Diffusers và Kernels, để người dùng có thể tải mô hình bằng các thiết lập như devicemap="tpu" thay vì chuyển code sang JAX hoặc phụ thuộc vào PyTorch/XLA. TPU là dòng chip do Google thiết kế, chuyên xử lý các phép tính ma trận được mạng nơ-ron sử dụng nhiều, và có thể thuê qua Google Cloud từ một chip riêng lẻ đến các cụm lớn. TorchTPU hỗ trợ PyTorch ở eager mode, huấn luyện phân tán, các kernel Pallas tối ưu cho TPU và torch.compile; với Diffusers, pipeline có thể chuyển sang TPU bằng pipe.to("tpu"), còn các mô hình lớn hơn có thể được chia trên nhiều chip bằng tensor parallelism. Hugging Face cho biết các ví dụ trên Cloud TPU v6e gồm những tác vụ tạo ảnh và video, trong đó compiled mode giúp rút ngắn đáng kể một số thời gian sinh nội dung; hãng cũng báo cáo mức tăng tốc giải mã tối đa 21,5 lần ở batch size 64 cho Qwen3-30B-A3B khi dùng kernel. Về thực tế, dự án này giảm đáng kể rào cản phần mềm cho các lập trình viên PyTorch muốn huấn luyện hoặc phục vụ mô hình Hugging Face trên TPU Google Cloud, dù người dùng vẫn cần quyền truy cập vào phần cứng TPU hoặc một TPU VM.
Lập trình viên đang dùng PyTorch và Google Cloud TPU có thể thử TorchTPU để chạy Transformers, Diffusers và các mô hình Hugging Face bằng code gần như PyTorch thông thường, thay vì phải chuyển sang JAX. Các đội triển khai mô hình lớn nên chú ý đến `torch.compile`, tensor parallelism và các TPU kernel để tận dụng nhiều chip hơn.
Đọc bài gốc