TPU 8: Google chia tách chip huấn luyện và suy luận, loại bỏ x86 để giành lợi thế AI

Google vừa công bố hai bộ tăng tốc AI mới do chính hãng phát triển tại hội nghị Cloud Next thường niên ở Las Vegas: một chip được thiết kế để tăng tốc độ huấn luyện và chip còn lại nhằm giảm chi phí phục vụ mô hình.

"Ông lớn" tìm kiếm khẳng định rằng các đơn vị xử lý tensor thế hệ thứ tám (TPU 8) của mình nhanh hơn tới 2,8 lần trong việc huấn luyện và mang lại hiệu suất trên mỗi đô la cao hơn 80% cho việc suy luận LLM (Large Language Model) so với TPU Ironwood của năm ngoái.

Để đạt được điều này, Google đã chia đôi lộ trình phát triển bộ tăng tốc, xây dựng TPU 8t cho huấn luyện và TPU 8i cho suy luận. Mặc dù các chip này được xây dựng trên nền tảng tương tự, mỗi loại đều nhắm đến việc loại bỏ các nút thắt cổ chai trong khối lượng công việc tương ứng.

Google không phải là công ty đầu tiên đi theo con đường này. Amazon Web Services (AWS) sớm nhận ra nhu cầu về các bộ tăng tốc được tối ưu hóa cho suy luận và huấn luyện trong giai đoạn đầu phát triển chip AI. Nvidia cũng đã thử nghiệm với loại chuyên môn hóa này, dù không ở cùng mức độ, khi thế hệ Blackwell Ultra được tối ưu hóa cụ thể cho suy luận AI.

Bên cạnh việc chia đôi lộ trình, Google đang loại bỏ bộ xử lý x86 để thay thế bằng CPU Axion dựa trên kiến trúc Arm do tự phát triển làm máy chủ cho TPU. Amazon đã thực hiện điều tương tự với Graviton và Trainium 3 vào đầu năm nay.

TPU 8t: Tối ưu cho quy mô huấn luyện khổng lồ

Giống như các bộ tăng tốc trước đây, TPU 8t của Google được thiết kế để huấn luyện ở quy mô cực lớn. Dưới "nắp capo", Google đã điều chỉnh hỗn hợp của các bộ tăng tốc vector, nhân ma trận và SparseCore được giới thiệu với chip Ironwood để tối đa hóa thông lượng dấu chấm động hiệu quả.

Mỗi bộ tăng tốc có 216 GB bộ nhớ HBM (High-Bandwidth Memory) với băng thông 6,5 TB/s, 128 MB SRAM trên chip, hiệu suất tính toán dấu chấm động 4-bit lên tới 12,6 petaFLOPS và băng thông chip-to-chip lên tới 19,2 Tbps.

So với GPU Rubin của Nvidia, TPU 8 mới có vẻ khiêm tốn hơn một chút về thông số kỹ thuật cá nhân. Rubin tự hào có hiệu suất huấn luyện FP4 lên tới 35 petaFLOPS và 288 GB HBM4 với băng thông 22 TB/s. Tuy nhiên, Google có lợi thế về khả năng mở rộng.

GPU mới nhất của Nvidia hỗ trợ tối đa 576 bộ tăng tốc trong một miền NVLink duy nhất trước khi phải mở rộng qua Ethernet hoặc InfiniBand. Ngược lại, TPU 8t sử dụng công tắc mạch quang (optical-circuit switches) để kết nối tới 9.600 bộ tăng tốc trong một nhóm thống nhất (pod). Nhiều pod sau đó được nối với nhau bằng mạng Virgo mới để hỗ trợ các miền tính toán lớn hơn.

Google cũng phát triển hệ thống lưu trữ Managed Lustre có khả năng cung cấp 10 TB/s dữ liệu tổng hợp trực tiếp vào bộ nhớ của bộ tăng tốc. Kết hợp với các khả năng độ tin cậy cải tiến, Google cho phép các TPU tối ưu cho huấn luyện hoạt động với "goodput" (hiệu suất thực tế) là 97%.

"Mỗi lỗi phần cứng, sự cố mạng hoặc khởi động lại checkpoint là thời gian mà cụm máy không huấn luyện, và ở quy mô huấn luyện tiên phong, mỗi điểm phần trăm có thể chuyển thành nhiều ngày thời gian huấn luyện tích cực," công ty giải thích trong một bài đăng trên blog.

TPU 8i: Điểm chuyển mình cho suy luận AI

Suy luận là khối lượng công việc tự hồi quy, nghĩa là với mỗi token được tạo ra, toàn bộ trọng số hoạt động của mô hình cần được truyền qua bộ nhớ. Do đó, nút thắt cổ chai chính thường là băng thông bộ nhớ.

Với TPU 8i tập trung vào suy luận, Google đánh đổi một phần sức mạnh tính toán (FLOPS) để lấy bộ nhớ đệm SRAM lớn hơn nhiều và vùng nhớ nhanh hơn, dung lượng cao hơn.

Chip này có thông số kỹ thuật tương đương với bộ tăng tốc Blackwell của Nvidia. TPU 8i có hiệu suất tính toán FP4 là 10,1 petaFLOPS, được hỗ trợ bởi 384 MB SRAM trên chip và 288 GB HBM với băng thông 8,6 TB/s.

Google cho biết lượng SRAM lớn này giúp TPU 8i giữ nhiều bộ nhớ đệm key-value (bộ nhớ ngắn hạn của mô hình) hơn trên chip, giảm thời gian các lõi phải chờ dữ liệu. TPU 8i cũng loại bỏ SparseCore để thay thế bằng động cơ tăng tốc tập thể (Collective Acceleration Engine - CAE), giúp giảm độ trễ đồng bộ hóa.

Cấu trúc mạng mới có tên Boardfly cho phép kết nối 1.152 chip (1.024 chip hoạt động tại bất kỳ thời điểm nào) bằng các công tắc mạch quang. Lợi thế chính của Boardfly là cắt giảm độ trễ tối đa giữa các chip từ 16 bước nhảy (hops) trong cấu trúc 3D torus xuống chỉ còn bảy bước nhảy.

Cả hai bộ tăng tốc TPU 8 sẽ có mặt trên Google Cloud Platform vào cuối năm nay dưới dạng phiên bản (instance) hoặc là một phần của nền tảng AI Hypercomputer.