Google và AWS đang giúp các mô hình PyTorch chạy trực tiếp hơn trên TPU và Trainium, qua đó mở rộng lựa chọn ngoài GPU Nvidia và CUDA. TorchTPU của Google đã được mở mã nguồn, còn TorchNeuron của AWS bước vào giai đoạn thử nghiệm công khai. Tuy nhiên, chuyển mô hình sang chip mới vẫn đơn giản hơn việc tối ưu hiệu năng cho môi trường sản xuất.
Tại Hội nghị PyTorch Bắc Mỹ 2026, Google công bố mở mã nguồn TorchTPU và AWS giới thiệu bản thử nghiệm công khai của TorchNeuron. Theo The New Stack, hai công cụ đưa PyTorch lên chip AI của Google và AWS. Nhờ đó, nhà phát triển không còn phụ thuộc hoàn toàn vào GPU Nvidia và hệ sinh thái CUDA.
Thay đổi này mang lại thêm lựa chọn cho các nhóm đang xây dựng hạ tầng AI. Dù vậy, chạy được mô hình trên phần cứng mới và khai thác hiệu năng cao nhất vẫn là hai bài toán khác nhau.
Google đưa PyTorch đến gần TPU hơn
TorchTPU cho phép chạy các khối lượng công việc PyTorch hiện có trên TPU mà không cần JAX trong môi trường thực thi. Theo PyTorch, phần lớn quá trình chuyển đổi chỉ cần những thay đổi nhỏ trong mã nguồn. Khối lượng công việc đòi hỏi hiệu năng cao vẫn có thể cần tối ưu riêng.
Google cho một số khách hàng xem trước TorchTPU vào tháng 4. Theo PyTorch, Google đã phát triển dự án từ giữa năm 2025 nhằm đơn giản hóa việc chạy khối lượng công việc PyTorch trên TPU.
Google xây dựng TorchTPU cùng nhóm PyTorch tại Meta. Nhà phát triển có thể dùng chế độ thực thi tức thời để gỡ lỗi mô hình. Khi cần tăng mức sử dụng TPU và giảm phần bộ nhớ phát sinh, họ có thể dùng torch.compile.
Google cho biết việc mở mã nguồn TorchTPU nhằm làm cho hệ sinh thái phần mềm TPU cởi mở và dễ tiếp cận hơn với cộng đồng PyTorch. Mã nguồn của dự án cũng đã được đưa lên GitHub.
Cùng với TorchTPU, Google công bố các tích hợp với Hugging Face, vLLM, SGLang, llm-d, TorchTitan, Ray và TorchRec. Công ty cũng đóng góp LiteRT cho hệ sinh thái PyTorch. Các sản phẩm và tích hợp này hỗ trợ vận hành PyTorch trên Google Cloud hoặc hạ tầng riêng.
AWS áp dụng cách tiếp cận tương tự với Trainium
TorchNeuron đưa khả năng làm việc trực tiếp với PyTorch lên Trainium, dòng chip AI của AWS dành cho huấn luyện và suy luận học sâu. Bản thử nghiệm công khai hỗ trợ Trainium2 và Trainium3, đồng thời cho phép thử một số mô hình bằng cách thay đổi cấu hình thay vì viết lại mã nguồn.
Theo AWS, ở chế độ thực thi tức thời, nhà phát triển có thể chuyển mô hình sang Trainium bằng cách đổi thiết bị với lệnh .to('neuron'). Với mô hình được biên dịch, họ phải chọn Neuron làm phần xử lý mặc định hoặc chỉ định Neuron trong lời gọi torch.compile.
AWS đã công bố TorchNeuron cho Trainium dưới dạng mã nguồn mở. Mã nguồn của công cụ cũng được cung cấp trên GitHub.
TorchNeuron tích hợp với Giao diện nhân Neuron, một ngôn ngữ và trình biên dịch cấp thấp dùng để lập trình các thiết bị Neuron thuộc Trainium và Inferentia. Nhà phát triển có thể tạo toán tử riêng và tối ưu hiệu năng trên Trainium. AWS và PyTorch cũng cung cấp hướng dẫn về các chế độ thực thi và sự đánh đổi hiệu năng của phần xử lý này.
Chuyển mô hình dễ hơn chưa đồng nghĩa hiệu năng tối đa
TorchTPU và TorchNeuron giảm lượng công việc ban đầu khi chuyển mô hình PyTorch sang phần cứng mới, nhưng không loại bỏ nhu cầu tối ưu cho từng loại chip. Mức hiệu năng cuối cùng còn phụ thuộc vào trình biên dịch, cơ chế truyền thông và các công cụ riêng của TPU hoặc Trainium.
Đây cũng là một yếu tố trong cuộc tranh luận về hệ sinh thái AI mở và đóng của Nvidia. PyTorch cho biết một thay đổi mã nguồn đơn giản có thể giúp mô hình chạy trên TPU, nhưng nhóm phát triển vẫn có thể phải tối ưu riêng để đạt hiệu năng cao nhất.
Với khối lượng công việc phân tán, PyTorch cho biết cách Skinny Collectives mặc định của TorchTPU có thể làm hiệu năng đầu cuối thấp hơn khoảng 30 đến 40% so với cách đã được tối ưu. Nguyên nhân là cách mặc định ngăn trình biên dịch XLA chồng lấp quá trình tính toán và truyền thông.
Nội dung tối ưu cho môi trường sản xuất
Đối với hệ thống sản xuất cần hiệu năng cao nhất, PyTorch khuyến nghị tinh chỉnh thêm bằng nhân Pallas. Các ngăn xếp được đề xuất gồm vLLM-TorchTPU, SGLang-TorchTPU và TorchTitan-TorchTPU.
TorchNeuron vẫn trong giai đoạn thử nghiệm công khai. AWS cho biết trong phần lớn trường hợp, các nhóm có thể thử Trainium bằng cách thay đổi cấu hình thay vì viết lại mã nguồn hoặc mô hình. Một số khối lượng công việc vẫn có thể cần nhiều thay đổi hơn.
Một nghiên cứu tình huống của AWS mô tả cách hãng điều chỉnh mô hình Laplace 3B cho Trainium2 cùng công ty AI i14. Việc triển khai cần ba thay đổi: dùng PyTorch nguyên bản cho nhân chú ý GLA, thiết kế lại cách Muon giao tiếp và biên dịch lượt lan truyền ngược tuần tự.
Lựa chọn mới cho nhóm phát triển PyTorch
TorchTPU và TorchNeuron chưa làm mất lợi thế hiện có của CUDA hoặc GPU Nvidia. Hai công cụ giúp mã nguồn PyTorch tiếp cận TPU và Trainium bằng con đường trực tiếp hơn, qua đó giảm công việc ban đầu khi các nhóm kỹ thuật đánh giá phần cứng mới.
Giới hạn hiện tại nằm ở giai đoạn tối ưu. Một mô hình có thể chạy sau vài thay đổi nhỏ, nhưng để đạt hiệu năng phù hợp cho môi trường sản xuất, nhóm kỹ thuật vẫn phải hiểu trình biên dịch, cơ chế truyền thông và công cụ riêng của từng nền tảng.
Hãy theo dõi Sine để cập nhật cách các nền tảng phần cứng mới đang thay đổi lựa chọn triển khai PyTorch và hạ tầng AI.
