TorchTPU: Google mang trải nghiệm PyTorch gốc lên phần cứng TPU với hiệu suất tối ưu

TorchTPU: Google mang trải nghiệm PyTorch gốc lên phần cứng TPU với hiệu suất tối ưu

TorchTPU Architecture

Bối cảnh xây dựng cơ sở hạ tầng cho AI hiện đại đã thay đổi căn bản. Biên giới mới của học máy đòi hỏi việc khai thác các hệ thống phân tán trải dài trên hàng nghìn bộ tăng tốc. Khi các mô hình mở rộng quy mô để chạy trên các cụm chip lên tới hàng trăm nghìn đơn vị, phần mềm điều khiển chúng phải đáp ứng những yêu cầu mới về hiệu suất, tính di động phần cứng và độ tin cậy.

Tại Google, các đơn vị xử lý tensor (TPU) là nền tảng của hạ tầng siêu máy tính của chúng tôi. Các mạch ASIC tùy chỉnh này cung cấp năng lượng cho việc huấn luyện và phục vụ các nền tảng AI của Google như Gemini và Veo, cũng như khối lượng công việc khổng lồ của khách hàng Cloud. Cả cộng đồng AI đều cần khả năng tiếp cận đầy đủ sức mạnh của TPU, và vì nhiều người dùng trong số này xây dựng mô hình bằng PyTorch, việc tích hợp để PyTorch hoạt động tự nhiên và hiệu quả trên TPU là vô cùng quan trọng.

Đó chính là lý do TorchTPU ra đời. Với tư cách là một đội ngũ kỹ thuật, nhiệm vụ của chúng tôi là xây dựng một stack phần mềm ưu tiên tính dễ sử dụng, khả năng chuyển đổi và hiệu suất xuất sắc. Chúng tôi muốn cho phép các nhà phát triển di chuyển khối lượng công việc PyTorch hiện có với ít thay đổi mã nhất có thể, đồng thời cung cấp các API và công cụ để khai thác tối đa sức mạnh tính toán từ phần cứng.

Hiểu về phần cứng TPU

Để hiểu về TorchTPU, trước hết cần hiểu phần cứng mà nó nhắm tới. Một hệ thống TPU không chỉ là một con chip; đó là một mạng lưới tích hợp. Một máy chủ được kết nối với nhiều chip, và mỗi chip kết nối với máy chủ cũng như các chip khác thông qua kết nối liên chip (ICI) của chúng tôi. ICI này liên kết các chip thành một tô pô hình xuyến 2D hoặc 3D hiệu quả cao, cho phép mở rộng quy mô lớn mà không gặp các nút thắt mạng truyền thống.

Bên trong mỗi chip, việc thực thi được chia sẻ giữa TensorCores và SparseCores. TensorCores là các đơn vị luồng đơn dành riêng cho toán học ma trận dày đặc, trong khi SparseCores xử lý các mẫu truy cập bộ nhớ không đều như nhúng (embeddings), các thao tác gather/scatter và giảm tải tập hợp (collectives).

TPU Infrastructure

Triết lý "Eager First" và các chế độ thực thi

Nguyên tắc cốt lõi của chúng tôi về tính dễ sử dụng rất đơn giản: nó phải cảm thấy như PyTorch. Một nhà phát triển có thể lấy một tập lệnh PyTorch hiện có, thay đổi khởi tạo thành "tpu" và chạy vòng lặp huấn luyện mà không cần sửa đổi một dòng logic nào.

Để đạt được điều này, chúng tôi đã thiết lập triết lý "Eager First" (Ưu tiên Eager). Thay vì yêu cầu nhà phát triển biên dịch đồ thị tĩnh ngay lập tức, chúng tôi triển khai TorchTPU bằng cách sử dụng giao diện "PrivateUse1" của PyTorch. Không có lớp con, không có bao bọc (wrapper); chỉ là các Tensors PyTorch quen thuộc trên TPU.

Chúng tôi đã thiết kế ba chế độ eager riêng biệt để hỗ trợ vòng đời phát triển:

1. Debug Eager: Phân phối một thao tác tại một thời điểm và đồng bộ hóa với CPU sau mỗi lần thực thi. Nó vốn chậm, nhưng vô giá để theo dõi sự không khớp về hình dạng, giá trị NaN và sự cố hết bộ nhớ.
2. Strict Eager: Duy trì phân tích đơn thao tác nhưng thực thi không đồng bộ, nhằm phản ánh trải nghiệm PyTorch mặc định. Điều này cho phép cả CPU và TPU thực thi đồng thời cho đến khi đạt được điểm đồng bộ hóa trong tập lệnh của người dùng.
3. Fused Eager: Đây là bước đột phá. Sử dụng phản chiếu tự động trên luồng thao tác, TorchTPU sẽ hợp nhất các bước trên đường bay thành các khối tính toán dày đặc lớn hơn trước khi chuyển chúng cho TPU. Bằng cách tối đa hóa việc sử dụng TensorCore và giảm thiểu chi phí băng thông bộ nhớ, Fused Eager luôn mang lại hiệu suất tăng từ 50% đến hơn 100% so với Strict Eager mà không cần người dùng thiết lập gì thêm.

Tất cả ba chế độ đều được hỗ trợ bởi Bộ nhớ đệm Biên dịch chia sẻ có thể hoạt động trên một máy chủ duy nhất hoặc được cấu hình liên tục trên các thiết lập đa máy chủ.

Biên dịch và Tối ưu hóa với XLA

Đối với những người dùng muốn mở khóa hiệu suất đỉnh cao trên TPU, TorchTPU tích hợp nguyên bản với giao diện torch.compile để biên dịch đồ thị đầy đủ. Chúng tôi bắt đầu bằng cách chụp đồ thị FX bằng Torch Dynamo. Tuy nhiên, thay vì định tuyến qua Torch Inductor, chúng tôi sử dụng XLA làm trình biên dịch backend chính.

Đây là một quyết định kiến trúc có chủ đích. XLA đã được kiểm chứng kỹ lưỡng cho các tô pô TPU. Quan trọng hơn, nó hiểu rõ cách tối ưu hóa sự chồng chéo quan trọng giữa tính toán dày đặc và truyền thông tập hợp qua ICI. Lớp dịch của chúng tôi ánh xạ các toán tử của PyTorch trực tiếp sang StableHLO, biểu diễn trung gian (IR) chính của XLA cho toán học tensor. Điều này tạo ra kết nối trực tiếp từ PyTorch vào đường dẫn hạ cấp (lowering path) cốt lõi của XLA.

Hỗ trợ phân tán và tùy chỉnh

Để duy trì tính linh hoạt ở quy mô lớn, TorchTPU tập trung mạnh vào các API phân tán của PyTorch. Hiện tại, TorchTPU hỗ trợ Distributed Data Parallel (DDP), Fully Sharded Data Parallel v2 (FSDPv2) và DTensor của PyTorch ngay lập tức.

Một hạn chế lớn của PyTorch/XLA (tiền thân của TorchTPU) là nó chỉ hỗ trợ mã SPMD thuần túy. Thực tế là các đầu vào của PyTorch thường có sự phân kỳ nhỏ trong mã chạy trên các thứ hạng (rank) khác nhau. TorchTPU được kiến trúc để hỗ trợ cẩn thận các thực thi phân kỳ (MPMD) và cô lập các nguyên thủy truyền thông khi cần thiết để bảo toàn tính đúng đắn với chi phí tối thiểu.

Lộ trình phát triển tới năm 2026

Chúng tôi đã đặt nền móng vững chắc cho việc hỗ trợ huấn luyện và phục vụ, và đang tích cực giải quyết một số thách thức mở. Một trọng tâm chính là giảm các lần biên dịch lại được kích hoạt bởi độ dài chuỗi và kích thước lô (batch size) động. Bằng cách triển khai tính động có giới hạn (bounded dynamism) nâng cao trong XLA, chúng tôi nhằm xử lý các thay đổi hình dạng mà không gây ra chi phí biên dịch.

Ngoài ra, chúng tôi đang xây dựng một thư viện toàn diện các nhân TPU được biên dịch trước cho các thao tác chuẩn để giảm đáng kể độ trễ của lần lặp thực thi đầu tiên. TorchTPU đại diện cho nỗ lực kỹ thuật của chúng tôi nhằm cung cấp trải nghiệm PyTorch liền mạch, hiệu suất cao trên phần cứng TPU.