TorchTPU đưa PyTorch chạy trực tiếp trên TPU của Google

TPU trở thành thiết bị PyTorch gốc

Google và nhóm PyTorch tại Meta đã mã nguồn mở TorchTPU, backend PyTorch gốc cho TPU, tại hội nghị PyTorch Conference 2026 ở San Jose vào ngày 20/10/2026. Với TorchTPU, TPU đứng cùng GPU và CPU như một lựa chọn thiết bị gốc trong PyTorch. Mã PyTorch hiện có có thể nhắm tới TPU thông qua thiết bị chuẩn torch.device("tpu"), và trong nhiều trường hợp việc chuyển đổi được mô tả là đổi lệnh thiết bị từ .to("cuda") sang .to("tpu").

Sự thay đổi này quan trọng vì trước đây con đường chính để chạy PyTorch trên TPU là PyTorch/XLA. PyTorch/XLA ghi lại đồ thị tensor lười, yêu cầu nhà phát triển đánh dấu vùng có thể biên dịch và chỉ hỗ trợ mã SPMD thuần túy. TorchTPU thay thế PyTorch/XLA và được Google khuyến nghị cho mọi khối lượng công việc mới, theo các bản tin về thông báo của Google.

Cách TorchTPU hoạt động

TorchTPU cắm vào PyTorch thông qua giao diện PrivateUse1. Tensor trong TorchTPU là tensor PyTorch thông thường, không có lớp con hay lớp bao bọc, và hệ thống chạy một tiến trình trên mỗi thiết bị trong khi vẫn giữ các tensor view và phép biến đổi tại chỗ. TorchTPU cũng xử lý thực thi MPMD phân kỳ, dạng thực thi mà PyTorch/XLA trước đây không hỗ trợ đầy đủ.

Quy trình làm việc quen thuộc của PyTorch được giữ nguyên. Nhà phát triển gỡ lỗi ở chế độ eager, rồi chuyển mô hình sang torch.compile để đạt thông lượng ở mức sản xuất. Hệ thống hỗ trợ khoảng 450 phép toán ATen, dùng kernel Pallas và các phép toán phân tán tích hợp sẵn, thay vì gắn thêm sau. TorchTPU được phát hành dưới giấy phép Apache 2.0, phiên bản 0.2.0 trên PyPI, yêu cầu Python 3.11 trở lên.

Hệ sinh thái đi kèm

Bản phát hành gồm 4 kho mã nguồn: TorchTPU, vLLM-TorchTPU, SGLang-TorchTPU và TorchTitan. Việc có backend gốc cho hai engine phục vụ suy luận phổ biến nhất là vLLM và SGLang giúp các mô hình đã chạy trên GPU có đường chuyển sang TPU mà không phải thay toàn bộ ngăn xếp phục vụ.

Google cũng cho biết đã hợp tác với Hugging Face để các mô hình trọng số mở trên Hugging Face Hub chạy gốc trên TorchTPU, và trong Transformers v5 việc triển khai mô hình lên TPU chỉ cần truyền device_map="tpu".

Về hiệu năng, Google báo cáo mức sử dụng FLOPs mô hình 52% khi huấn luyện Qwen 3 8B bằng TorchTitan trên Ironwood, với khả năng mở rộng gần tuyến tính lên tới 8.192 chip Ironwood. Đây là con số do Google tự báo cáo, và vẫn cần đặt cạnh thực tế rằng đổi thiết bị chỉ là điểm bắt đầu. Một trường hợp tinh chỉnh được các bản tin nêu ra cho thấy chế độ eager chưa tối ưu có thể chậm hơn nhiều so với bản đã tinh chỉnh đầy đủ, vì vậy hiệu năng đỉnh vẫn đòi hỏi điều chỉnh hình dạng tensor và phân mảnh.

Vì sao Meta quan tâm tới TPU

Meta là bên ủng hộ rõ ràng nhất cho TorchTPU, và lý do nằm ở chiến lược phần cứng. Giám đốc kỹ thuật PyTorch tại Meta, Jana van Greunen, được dẫn lời trong bài đăng của Google rằng Meta không thể bị ràng buộc bởi một kiến trúc phần cứng duy nhất. Theo bài đăng này, TorchTPU hỗ trợ các khối lượng công việc mô hình ngôn ngữ lớn, quảng cáo và gợi ý tại Meta, và cho phép công ty triển khai khối lượng công việc PyTorch lên TPU của Google.

Lập luận này vượt ra ngoài một công ty. Khi mã mô hình tách khỏi một loại chip cụ thể, các đội có thể thử nghiệm cụm TPU trên đám mây trong khi vẫn giữ một cơ sở mã duy nhất, rồi so sánh chi phí, thông lượng và mức tiêu thụ năng lượng với GPU. Trong bối cảnh suy luận chiếm phần lớn khối lượng tính toán của trung tâm dữ liệu, khả năng chuyển đổi phần cứng mà không viết lại mô hình trở thành một lợi thế vận hành thực sự.

Điều cần kiểm chứng trong thực tế

TorchTPU đã được định vị là sẵn sàng cho sản xuất, và Google cho biết một số khách hàng đám mây đã bắt đầu áp dụng backend này. Thước đo thực tế sẽ là số mô hình Hugging Face chạy được mà không gặp lỗi, mức độ ngang bằng số học so với thực thi trên GPU, và chi phí suy luận trên mỗi triệu token trong các khối lượng công việc thật. Các bản tin dẫn con số 98% mô hình Tier 1 không gặp lỗi và 97% ngang bằng số học so với GPU, nhưng đây vẫn là các con số do nhà cung cấp báo cáo và cần được kiểm chứng độc lập qua các đợt triển khai.

Nguồn tham khảo:

MarkTechPost

WebProNews