PyTorch trên TPU và Trainium: thoát phụ thuộc Nvidia đến đâu?

Chạy PyTorch trên TPU và Trainium nay chỉ cần đổi thiết bị nhờ TorchTPU, TorchNeuron. Bài phân tích số liệu, giới hạn, chi phí và ai nên thử trước.

Nội dung chính

Chạy PyTorch trên TPU của Google hay chip Trainium của AWS trước đây cần thư viện trung gian và khá nhiều chỉnh sửa. Ngày 19–20/10/2026, ngay trước PyTorch Conference, hai hãng cùng mở mã nguồn một "backend" PyTorch gốc: TorchTPU (Google, phát triển cùng đội PyTorch của Meta) và TorchNeuron (AWS, bản beta công khai). Mục tiêu chung: mã PyTorch viết cho GPU Nvidia chạy được trên chip của họ với rất ít thay đổi.

• Đổi thiết bị thay vì viết lại. Với TorchNeuron, AWS nói thay đổi duy nhất là model.to('neuron'). Với TorchTPU, Google nói chỉ cần "thay đổi mã tối thiểu"; trong Hugging Face Transformers v5 chỉ cần device_map="tpu".
• Đủ hai chế độ chạy. Cả hai hỗ trợ chế độ eager (chạy từng lệnh, dễ gỡ lỗi) và torch.compile (biên dịch cả đồ thị để chạy nhanh khi đưa vào sản xuất).
• Huấn luyện phân tán sẵn có. DDP, FSDP2, DTensor được hỗ trợ ở cả hai; TorchNeuron thêm song song tensor và song song pipeline.
• Gắn với hệ sinh thái phục vụ mô hình. TorchTPU có backend cho vLLM và SGLang; TorchNeuron có vLLM Neuron và hỗ trợ Transformers v5, Diffusers, trl.
• Mã nguồn mở Apache 2.0 cho cả hai dự án trên GitHub.

Phân tích

Phân tích: một dòng code có đủ để rời GPU Nvidia?

Câu trả lời ngắn: đủ để thử, chưa đủ để chuyển hẳn. Thay đổi quan trọng nằm ở cách tích hợp, còn chi phí thực tế phụ thuộc nhiều vào loại mô hình bạn chạy.

So với cách làm trước đây thì khác gì?

Điểm khác lớn nhất là PrivateUse1. Đây là "cổng cắm" chính thức của PyTorch cho phần cứng mới. Tensor trên TPU hay Trainium là tensor PyTorch bình thường, không bị bọc lại. Nhờ vậy các thư viện như Hugging Face, vLLM hay TorchTitan không phải viết nhánh riêng cho từng chip.

Vì sao chạy được mà không viết lại?

Khi bạn gọi torch.compile, PyTorch dùng TorchDynamo để chụp đồ thị tính toán. Thay vì chuyển đồ thị sang trình biên dịch dành cho GPU, TorchTPU chuyển nó sang StableHLO rồi để XLA tối ưu cho TPU. TorchNeuron đi đường tương tự qua Torch-MLIR, StableHLO, rồi trình biên dịch Neuron. Lập trình viên vẫn viết PyTorch quen thuộc; phần dịch sang chip nằm ở tầng dưới.

Chế độ "Fused Eager" của TorchTPU tự gom nhiều lệnh nhỏ thành khối lớn trước khi gửi xuống chip. Google từng công bố mức tăng 50% đến hơn 100% so với chạy từng lệnh. Đây là con số của hãng, chưa có đo độc lập.

Tác động tới từng nhóm

• Lập trình viên ML: có thể chạy cùng một đoạn huấn luyện trên GPU, TPU và Trainium để so sánh. Việc tự viết kernel vẫn cần học thêm Pallas (TPU) hoặc NKI (Trainium).
• Nhà quản lý kỹ thuật: có thêm đòn bẩy khi đàm phán giá GPU. Một hợp đồng đám mây không còn buộc chặt vào CUDA nếu mô hình chạy được trên nhiều chip.
• Doanh nghiệp nhỏ và vừa, kể cả ở Việt Nam: phần lớn dùng mô hình qua API nên ảnh hưởng gián tiếp. Nhóm tự phục vụ mô hình mở (Qwen, Llama) trên Google Cloud hoặc AWS là nhóm hưởng lợi trực tiếp.
• Marketer, đội sản phẩm: chưa có thay đổi ngay. Lợi ích nếu có sẽ đến qua giá API rẻ hơn từ nhà cung cấp dùng chip này.

Chi phí thực tế là bao nhiêu?

Hai hãng không công bố bảng giá mới. Con số chi phí duy nhất là của Deepgram: xử lý nhiều hơn khoảng 60% lượng âm thanh trên mỗi USD so với GPU họ đang dùng. Bảng dưới là ví dụ giả định, áp con số đó cho hai quy mô và giả định chi phí chuyển đổi 30.000 USD (khoảng 2 kỹ sư trong 1–2 tháng).

Kết luận từ bảng: với hóa đơn nhỏ, chuyển chip hiếm khi đáng công. Với hóa đơn lớn và mô hình ổn định, bài toán bắt đầu có lời. Mức 60% là kết quả của một khách hàng, một loại tác vụ (âm thanh); mô hình của bạn có thể cao hơn, thấp hơn, hoặc không có lợi gì.

Rủi ro và giới hạn

• Độ phủ toán tử chưa đủ. Trong lần chạy bộ kiểm tra PyTorch tháng 8/2026, khoảng 22.000 bài kiểm tra bị bỏ qua vì TorchNeuron chưa có toán tử tương ứng. Bản đầu chưa hỗ trợ hình dạng động (dynamic shapes), sparse tensor, NestedTensor; FP8 mới ở beta kín.
• Hình dạng động vẫn là điểm yếu của TPU. Google tự nhận độ dài chuỗi hoặc kích thước batch thay đổi sẽ khiến mô hình phải biên dịch lại, và lần chạy đầu có độ trễ cao. Chatbot với độ dài câu hỏi thay đổi liên tục là đúng loại tải bị ảnh hưởng.
• Khóa nhà cung cấp chỉ đổi chỗ. TPU chỉ có trên Google Cloud, Trainium chỉ có trên AWS. Thoát CUDA có thể đồng nghĩa với gắn chặt hơn vào một đám mây.
• Mã chỉ dành cho CUDA không chạy. CUDA graphs, các lệnh NCCL riêng và kernel CUDA tự viết vẫn phải viết lại.

Điểm cần kiểm chứng

• Mức 52% MFU và 8.192 chip là thử nghiệm nội bộ của Google trên TorchTitan, chưa có đối chiếu độc lập với GPU cùng mức giá.
• Con số Zenteiq (gấp khoảng 4 lần) và Deepgram (+60%/USD) là trích dẫn khách hàng do hãng chọn đăng.
• Con số "85% ứng dụng AI dùng PyTorch" mà TechCrunch dẫn không nêu nguồn cụ thể; TBN Grow chưa xác minh được.
• Cách tự kiểm tra: chọn một mô hình đang chạy thật, đo token/giây và chi phí/1 triệu token trên GPU hiện tại, rồi chạy cùng mô hình trên TPU v6e hoặc Trainium2 với cùng độ dài đầu vào.

Nhận định của TBN Grow

Đây là thay đổi về hạ tầng phần mềm, có giá trị thực cho đội tự huấn luyện hoặc tự phục vụ mô hình mở trên Google Cloud hay AWS. Đội có hóa đơn GPU từ vài chục nghìn USD/tháng nên thử ngay với một mô hình ổn định. Đội nhỏ, dùng API hoặc mô hình có đầu vào biến động mạnh nên chờ 1–2 bản phát hành tới, khi độ phủ toán tử và hỗ trợ hình dạng động hoàn thiện hơn.

Điều cần theo dõi

Điều cần theo dõi khi chạy PyTorch trên TPU và Trainium

Việc làm được ngay

• [ ] Kiểm tra mã của bạn có gọi trực tiếp cuda ở đâu; thay bằng API torch.accelerator để dễ chuyển chip.
• [ ] Thử TorchTPU: pip install torch_tpu==0.2.0 trên một máy TPU v6e (README của dự án ghi đây là cách dễ bắt đầu nhất; cần thêm --extra-index-url https://download.pytorch.org/whl/cpu để không cài bản CUDA).
• [ ] Thử TorchNeuron: pip install torch-neuronx trên máy trn2, đổi thiết bị sang neuron.
• [ ] Chạy bộ kiểm tra của chính dự án để tìm toán tử chưa được hỗ trợ trước khi tính đến chi phí.
• [ ] So sánh chi phí trên 1 triệu token hoặc trên 1 giờ âm thanh, không chỉ tốc độ.

Mốc thời gian và điều chưa công bố

• AWS hẹn đưa TorchAO và FP8 lên bản chính thức ở bản phát hành tới và lấp dần khoảng trống toán tử; chưa có ngày cụ thể.
• Google đặt mục tiêu hỗ trợ hình dạng động trong torch.compile và thư viện kernel biên dịch sẵn; chưa có mốc thời gian.
• Google chưa gắn nhãn "beta" hay "chính thức" cho TorchTPU và chưa công bố danh sách đầy đủ các thế hệ TPU được hỗ trợ.
• Chưa có phản hồi chính thức từ Nvidia.

Đọc thêm: Thị trường chip AI Trung Quốc đổi chiều: Nvidia còn 34%, Huawei lên 28% và bài toán hoàn vốn của đầu tư hạ tầng AI.

Nguồn kiểm chứng: Google Cloud Blog, Google Developers Blog, GitHub TorchTPU, PyTorch Blog – TorchNeuron, AWS What's New, TechCrunch.