Google mã nguồn mở TorchTPU: đưa PyTorch chạy gốc trên chip TPU

Google vừa chính thức mã nguồn mở TorchTPU, một back end gốc cho phép các mô hình PyTorch chạy trực tiếp trên chip Tensor Processing Unit (TPU) của hãng. Động thái này được giới phát triển AI đặc biệt quan tâm, bởi từ trước đến nay TPU vốn gắn chặt với hệ sinh thái JAX và TensorFlow, còn PyTorch — framework phổ biến nhất trong cộng đồng nghiên cứu — lại gần như "đóng khung" với GPU của NVIDIA.

Biểu đồ mức độ quan tâm dự án TorchTPU trên GitHub

TorchTPU mang đến điều gì?

Về bản chất, TorchTPU là một back end PyTorch dành riêng cho TPU, bao gồm các thành phần chính:

• Kernel ATen tùy chỉnh cho nhiều phép toán PyTorch, giúp các toán tử chạy hiệu quả trên kiến trúc TPU.
• Bộ nhớ đệm biên dịch (compilation cache) tối ưu việc thực thi các phép tính XLA, giảm thời gian biên dịch lặp lại.
• Hỗ trợ huấn luyện phân tán với Data Parallelism và Tensor Parallelism.
• Tích hợp torch.compile() để chạy ở chế độ graph, tận dụng tối đa khả năng tối ưu của XLA.
• Bộ công cụ debug và benchmark phục vụ quá trình phát triển và đo hiệu năng.

Đáng chú ý, kho mã nguồn đã có sẵn ví dụ cho nhiều mô hình phổ biến như Llama, Qwen, ResNet, DLRM, minGPT và UNet — tức là từ mô hình ngôn ngữ lớn cho tới thị giác máy tính và hệ thống gợi ý.

Cài đặt và trải nghiệm

TorchTPU hiện hỗ trợ Python 3.11 trở lên, theo đúng chính sách tương thích phiên bản của PyTorch. Người dùng có thể cài bản phát hành công khai mới nhất (0.2.0 trên PyPI) bằng pip, lưu ý cần thêm --extra-index-url https://download.pytorch.org/whl/cpu để tránh pip kéo nhầm bản torch dành cho CUDA.

Với những ai muốn tự build từ mã nguồn, dự án sử dụng Bazel để quản lý phụ thuộc và đóng gói wheel. Tài liệu hướng dẫn chi tiết cùng các notebook tương tác Marimo được công bố tại torchtpu.readthedocs.io — người dùng cần chạy trên một TPU VM, ví dụ instance v6e trên Google Cloud, để thử nghiệm.

Dự án được phát hành theo giấy phép Apache 2.0 và mở cửa cho cộng đồng đóng góp, kèm quy trình ký Contributor License Agreement (CLA) và review mã nguồn rõ ràng.

Ý nghĩa với cộng đồng AI

Việc TorchTPU xuất hiện cho thấy Google đang nỗ lực phá bỏ rào cản giữa PyTorch và phần cứng TPU. Với các nhóm nghiên cứu và doanh nghiệp đã đầu tư vào TPU trên Google Cloud nhưng bị "khóa" trong hệ sinh thái JAX, đây là lựa chọn mới giúp tái sử dụng mã nguồn PyTorch sẵn có mà không phải viết lại.

Ở góc độ cạnh tranh, động thái này được xem là nước đi nhằm thu hẹp khoảng cách với hệ sinh thái CUDA của NVIDIA — vốn đang thống trị thị trường huấn luyện AI. Tuy nhiên, TorchTPU vẫn ở giai đoạn đầu (phiên bản 0.2.0), nên độ chín và hiệu năng thực tế vẫn cần được cộng đồng kiểm chứng.

Với các lập trình viên Việt Nam làm việc trong lĩnh vực AI, dự án này mở ra cơ hội thử nghiệm mô hình trên hạ tầng TPU — vốn có thể tiếp cận qua Google Cloud — mà không cần chuyển đổi toàn bộ pipeline từ PyTorch sang JAX. Đây là tín hiệu tích cực cho tính linh hoạt của công cụ AI mã nguồn mở.