Google công khai mã nguồn TorchTPU, đưa PyTorch nguyên bản lên các TPU của mình

Ngày 20 tháng 10 năm 2026, tại PyTorch Conference 2026, Google thông báo TorchTPU — bộ công cụ PyTorch nguyên bản dành cho Tensor Processing Unit (TPU) — hiện đã được phát hành dưới dạng mã nguồn mở trên GitHub, theo thông tin công ty đăng trên blog Google Cloud. Phần mềm này được đồng phát triển với nhóm PyTorch của Meta, nhằm chạy các tác vụ PyTorch hiện có trên TPU mà chỉ cần thay đổi tối thiểu về mã.

Google cho biết với bản phát hành này, TPU trở thành lựa chọn nguyên bản bên cạnh GPU và CPU cho PyTorch, hỗ trợ huấn luyện, phục vụ mô hình và triển khai thông qua các công cụ như vLLM, SGLang, llm-d, TorchTitan, Ray, TorchRec, Hugging Face và LiteRT. Nhà phát triển có thể gỡ lỗi mô hình ở chế độ eager, sau đó chuyển sang hiệu năng cấp sản xuất thông qua torch.compile để tối ưu mức sử dụng TPU và giảm chi phí bộ nhớ.

Trong thông báo, Jana van Greunen, Giám đốc Kỹ thuật PyTorch tại Meta, cho biết TorchTPU cho phép Meta linh hoạt đồng phát triển phần cứng cùng nhóm hệ thống học máy nòng cốt, đồng thời hỗ trợ các tác vụ quy mô lớn của LLM, quảng cáo và hệ thống đề xuất. Bà nói: “Công nghệ này cho phép chúng tôi triển khai các tác vụ PyTorch trên TPU của Google, trực tiếp phản hồi cho các nhóm kỹ thuật và xây dựng hệ sinh thái linh hoạt mà mọi người cần.”

Tích hợp cho hệ sinh thái Hugging và phục vụ mô hình

Google hợp tác với Hugging Face để các mô hình có trọng số mở trên Hugging Face Hub có thể chạy nguyên bản trên TorchTPU. Trong Transformers v5, việc triển khai mô hình lên TPU được thực hiện bằng cách truyền device_map="tpu", và thư viện Diffusers chạy nguyên bản trên TPU, hỗ trợ chế độ eager, torch.compile và xử lý song song tensor trên nhiều bộ tăng tốc. Google cho biết họ cũng xây dựng các quy trình CI/CD chuyên biệt để kiểm thử liên tục hằng đêm bộ Hugging Face Transformers trên hạ tầng TPU. Jeff Boudier, Phó chủ tịch phụ trách Sản phẩm tại Hugging Face, cho biết sự hợp tác này hiện thực hóa cam kết của nền tảng về quyền lựa chọn phần cứng.

Đối với việc phục vụ mô hình, backend vLLM nguyên bản của PyTorch trên TorchTPU được phát triển với sự hợp tác của Inferact. Google cho biết giải pháp này hỗ trợ xử lý theo lô liên tục, lượng tử hóa và xử lý song song đa chip trên phần cứng TPU bằng các kernel Pallas tùy chỉnh, đồng thời giữ nguyên mô hình worker, bộ lập lịch và máy chủ API ở phiên bản upstream để các hệ thống suy luận hiện có không cần thay đổi. Simon Mo, Giám đốc điều hành kiêm đồng sáng lập Inferact, cho biết tích hợp này mang đến khả năng hỗ trợ vLLM hạng nhất trên TPU.

Backend PyTorch nguyên bản dành cho SGLang, được phát triển cùng RadixArk, hướng đến các tác vụ có cấu trúc và tác vụ dạng tác tử. Theo thông báo, SGLang sử dụng RadixAttention để tái sử dụng bộ nhớ đệm KV giữa các yêu cầu, nhờ đó những tiền tố lời nhắc lặp lại không cần tính toán lại; phiên bản TorchTPU sử dụng backend truyền tải phân tách, đồng thời làm backend I/O cho HiCache. Banghua Zhu, Giám đốc công nghệ kiêm đồng sáng lập RadixArk, cho biết TorchTPU đã đẩy nhanh công việc của công ty khi họ triển khai hỗ trợ SGLang trên nhiều loại chip khác nhau.

Google cũng phát hành llm-d 1.0 cùng Red Hat, IBM và cộng đồng gồm hơn 500 người đóng góp và 25 tổ chức, dưới dạng một dự án Sandbox của Cloud Native Computing Foundation. Thông báo mô tả llm-d là cầu nối giữa hệ sinh thái phục vụ mô hình PyTorch và hoạt động triển khai sản xuất dựa trên Kubernetes: bộ định tuyến sử dụng cơ chế định tuyến nhận biết bộ nhớ đệm tiền tố để chuyển từng yêu cầu đến bản sao đã lưu ngữ cảnh tương ứng, đồng thời tách bước prefill đòi hỏi nhiều năng lực tính toán khỏi bước decode tiêu tốn nhiều bộ nhớ. Phiên bản 1.0 bổ sung khả năng kiểm soát luồng đa người thuê, tính sẵn sàng cao cho bộ định tuyến, các API v1 ổn định và định tuyến nhận biết phiên cho tác tử; trên Google Kubernetes Engine, một điểm cuối llm-d dùng chung có thể chạy trên GPU hoặc TPU.

Trong thông báo, Giáo sư Sashikumaar Ganesan, nhà sáng lập Zenteiq, cho biết GKE Inference Gateway chạy trên llm-d đã mang lại thông lượng suy luận tổng hợp cao hơn 4 lần cho hoạt động tạo dữ liệu tổng hợp của công ty, nâng công suất tạo token mỗi ngày từ 1,47 tỷ lên 5,97 tỷ token trên mỗi 1.000 bộ tăng tốc TPU v6e.

Huấn luyện và điều phối

TorchTitan, bản triển khai tham chiếu nguyên bản cho PyTorch về huấn luyện phân tán, hiện có tích hợp TorchTPU mà Google mô tả là đạt tiêu chuẩn sản xuất, được xây dựng dựa trên torch.compile, SimpleFSDP của TorchTitan, các kernel Pallas chuyên dụng và các tối ưu hóa ở cấp đồ thị. TorchTitan cung cấp các công thức TorchTPU để tiền huấn luyện những mô hình đặc như Llama 3 70B và 405B, Qwen 3 8B và 32B, cùng MuseGlimmer 30B; ngoài ra còn có mô hình hỗn hợp chuyên gia Qwen 3 235B và transformer khuếch tán FLUX.1 12B. Google cho biết tích hợp này đã được xác thực trên tối đa 8.000 TPU Ironwood (TPU v7x), và hoạt động huấn luyện TorchTPU thông qua TorchTitan đạt mức sử dụng FLOPS của mô hình lên đến 52% với Qwen 3 8B, đồng thời mở rộng gần tuyến tính trên tối đa 8.192 TPU Ironwood.

Google cũng đóng góp một backend TPU và một công thức tham chiếu toàn trình cho VeRL, một khung học tăng cường. Công thức GRPO, được xây dựng trên kiến trúc trainer-sampler phân tách và trình diễn trên các bộ dữ liệu suy luận toán học, nhằm giúp nhà phát triển chạy các tác vụ VeRL trên TPU mà vẫn giữ nguyên mã mô hình và thuật toán hiện có, vốn không phụ thuộc vào phần cứng.

Ray bổ sung hỗ trợ TorchTPU nguyên bản để điều phối xử lý dữ liệu, huấn luyện, học tăng cường và phục vụ mô hình PyTorch trên CPU, GPU và TPU trong cùng một quy trình. Edward Oakes, kiến trúc sư trưởng tại Anyscale, cho biết hai công ty đã phối hợp để đưa TPU trở thành bộ tăng tốc hạng nhất trong Ray, đồng thời nói: “TorchTPU tích hợp trực tiếp TPU của Google vào ngăn xếp tiêu chuẩn ngành này.”

Kiến trúc trước đây và khả năng cung cấp

TorchTPU được giới thiệu công khai vào ngày 7 tháng 4 năm 2026 trong bài đăng trên blog Google Developers, mô tả thiết kế “Ưu tiên Eager” dựa trên giao diện PrivateUse1 của PyTorch, với các chế độ eager Debug, Strict và Fused sử dụng chung bộ nhớ đệm biên dịch. Khi đó, Google cho biết Fused Eager mang lại mức tăng hiệu năng từ 50% đến hơn 100% so với Strict Eager, đồng thời torch.compile trên TorchTPU thu thập các đồ thị FX bằng Torch Dynamo rồi hạ cấp chúng qua XLA và StableHLO; các kernel tùy chỉnh được viết bằng Pallas và JAX. Bài đăng cho biết ngăn xếp này hỗ trợ sẵn Distributed Data Parallel, FSDPv2 và DTensor, đồng thời được thiết kế để xử lý tình trạng thực thi khác nhau giữa các hạng. Bài đăng cũng liệt kê việc ra mắt kho lưu trữ GitHub công khai trong số các mục tiêu lộ trình năm 2026 của TorchTPU.

Kho lưu trữ torch_tpu hiện được công khai trong tổ chức google-pytorch trên GitHub và được cấp phép theo Apache 2.0. Tệp README liệt kê các kernel ATen tùy chỉnh, bộ nhớ đệm biên dịch cho các phép tính XLA, hỗ trợ huấn luyện phân tán theo kiểu song song dữ liệu và song song tensor, tích hợp torch.compile, các tiện ích gỡ lỗi và đo điểm chuẩn, cùng ví dụ cho Llama, Qwen, ResNet, DLRM, minGPT và UNet. Dự án hỗ trợ Python 3.11 trở lên; gói torch-tpu 0.2.0 được phát hành trên PyPI vào ngày 19 tháng 10 năm 2026, với các gói wheel Linux x86-64 dành cho CPython từ phiên bản 3.11 đến 3.14. Hướng dẫn sử dụng được đăng tại torchtpu.readthedocs.io; README lưu ý rằng các sổ tay tương tác của dự án phải được chạy từ máy ảo TPU, trong đó phiên bản GCE v6e là lựa chọn dễ dùng nhất.

Trong cùng thông báo, Google cho biết LiteRT, trước đây là TensorFlow Lite, đang gia nhập hệ sinh thái PyTorch: thư viện LiteRT-Torch chuyển đổi mô hình PyTorch chỉ bằng một litert-torch.convert() lệnh gọi, và các mô hình đã chuyển đổi có thể chạy trên thiết bị Android, iOS, macOS, Windows, Linux, web và IoT. Bộ công cụ tác tử TPU nhằm tự động hóa quá trình làm quen, chẩn đoán và tinh chỉnh hiệu năng suy luận cho các mô hình PyTorch hiện đang ở giai đoạn xem trước ban đầu trên GitHub.

Thông báo hướng dẫn nhà phát triển tham khảo tài liệu về TPU của Google Cloud để bắt đầu chạy PyTorch trên TPU.