Google tung TPU 8, tách riêng chip huấn luyện và suy luận để tăng tốc cuộc đua hạ tầng AI

Tại sự kiện Cloud Next ở Las Vegas, Google công bố hai bộ tăng tốc AI nội bộ thế hệ mới thuộc dòng TPU 8. Thay vì tiếp tục theo đuổi một kiến trúc chip dùng chung cho mọi tác vụ, hãng chia hẳn thành TPU 8t cho huấn luyện mô hình và TPU 8i cho suy luận. TPU, viết tắt của Tensor Processing Unit, là dòng chip chuyên dụng do Google tự phát triển để xử lý các phép toán ma trận và tensor vốn là nền tảng của học máy và AI hiện đại. Theo Google, TPU 8 có thể nhanh hơn tới 2,8 lần trong huấn luyện và mang lại hiệu năng trên mỗi đô la cao hơn 80% cho suy luận mô hình ngôn ngữ lớn, hay LLM, so với thế hệ Ironwood năm ngoái.

Một thay đổi đáng chú ý khác là Google loại bỏ bộ xử lý x86 khỏi vai trò máy chủ điều phối TPU và chuyển sang Axion, dòng CPU dựa trên kiến trúc Arm do hãng tự thiết kế. x86 là kiến trúc CPU phổ biến lâu nay trong máy chủ doanh nghiệp, chủ yếu gắn với Intel và AMD, còn Arm nổi tiếng nhờ hiệu quả điện năng tốt hơn và ngày càng được đẩy mạnh trong trung tâm dữ liệu. Động thái này cho thấy Google đang tối ưu toàn bộ ngăn xếp hạ tầng, từ CPU chủ, chip tăng tốc, đến mạng và lưu trữ, thay vì chỉ nâng cấp riêng bộ xử lý AI.

Phiên bản TPU 8t được tinh chỉnh cho huấn luyện, tức giai đoạn dạy mô hình học từ khối dữ liệu khổng lồ. Google cho biết mỗi chip có 216 GB HBM, tức bộ nhớ băng thông cao, đạt 6,5 TB/giây; 128 MB SRAM trên chip, là bộ nhớ cực nhanh đặt sát khối tính toán; cùng năng lực xử lý tới 12,6 petaFLOPS ở định dạng FP4. FP4 là kiểu số thực 4-bit, một định dạng độ chính xác thấp nhưng rất hữu ích trong AI vì giúp tăng tốc và giảm điện năng cho nhiều tác vụ. TPU 8t cũng có băng thông kết nối chip-to-chip lên tới 19,2 Tbps, cho thấy Google không chỉ tập trung vào sức mạnh từng chip mà còn nhắm tới khả năng ghép hàng nghìn chip thành một hệ thống thống nhất.

Nếu so từng đơn vị, TPU mới chưa chắc vượt các GPU cao cấp của Nvidia như Rubin, vốn được công bố có thể đạt 35 petaFLOPS FP4 và dùng bộ nhớ HBM4 với băng thông rất lớn. Nhưng trong huấn luyện các mô hình AI hàng đầu, doanh nghiệp không vận hành một chip đơn lẻ mà dùng hàng nghìn bộ tăng tốc cùng lúc. Đây là nơi Google tin mình có lợi thế. Hãng cho biết TPU 8t có thể kết nối tới 9.600 bộ tăng tốc trong một pod thống nhất bằng optical-circuit switch, tức công nghệ chuyển mạch quang theo mạch, hoạt động giống tổng đài điện thoại hơn là kiểu chuyển gói dữ liệu truyền thống. Cách làm này giúp giảm độ trễ và hạn chế tắc nghẽn khi hệ thống mở rộng quy mô. Phù hợp với cả máy tính công nghiệp, Embedded PC xử dụng tính toán biên AI hoặc các loại máy quét mã vạch AI và hệ thống Vision AI khác.

Sau cấp pod, Google dùng mạng mới mang tên Virgo Network để nối nhiều cụm lớn hơn nữa. Theo mô tả của hãng, kiến trúc mạng hai tầng gần như all-to-all, tức gần như mọi nút đều có thể giao tiếp trực tiếp với nhau qua ít tầng trung gian, cho phép hỗ trợ tới 134.000 TPU trong một trung tâm dữ liệu và thậm chí chạm mốc một triệu TPU khi liên kết nhiều địa điểm. Song song với đó là Managed Lustre, hệ thống lưu trữ hiệu năng cao có thể bơm dữ liệu vào bộ nhớ bộ tăng tốc với tổng băng thông 10 TB/giây. Lustre vốn là hệ thống tệp quen thuộc trong HPC, tức điện toán hiệu năng cao, nơi các siêu máy tính cần truy cập dữ liệu cực nhanh cho các bài toán lớn.

Một điểm thú vị trong thông điệp của Google là hãng nhấn mạnh chỉ số goodput đạt 97% cho TPU tối ưu huấn luyện. Khác với throughput, vốn chỉ đo lượng công việc lý thuyết mà hệ thống có thể xử lý, goodput phản ánh phần thời gian cụm máy thực sự làm việc có ích, ở đây là thực sự huấn luyện mô hình. Điều này rất quan trọng vì ở quy mô hàng chục nghìn chip, chỉ một lỗi phần cứng, một lần nghẽn mạng hay một lần khởi động lại checkpoint cũng có thể làm mất nhiều ngày huấn luyện. Checkpoint là các mốc lưu trạng thái mô hình để có thể tiếp tục từ giữa chừng nếu xảy ra sự cố.

Nếu TPU 8t phục vụ giai đoạn tạo ra mô hình, TPU 8i lại nhắm vào suy luận, tức giai đoạn mô hình trả lời người dùng sau khi đã được huấn luyện xong. Với các LLM, suy luận là bài toán tự hồi quy, hay auto-regressive, nghĩa là mô hình tạo từng token một và mỗi token mới đều cần truy xuất lại phần lớn trọng số đang hoạt động của mô hình. Token có thể hiểu đơn giản là đơn vị văn bản nhỏ, như một từ hoặc một phần của từ. Vì vậy, nút thắt không chỉ nằm ở năng lực tính toán mà còn ở tốc độ cấp dữ liệu từ bộ nhớ. TPU 8i giảm bớt FLOPS để đổi lấy 384 MB SRAM trên chip và 288 GB HBM với băng thông 8,6 TB/giây, nhằm giữ được nhiều dữ liệu hơn ngay gần lõi xử lý.

Google cho biết lượng SRAM lớn hơn giúp TPU 8i giữ được nhiều key-value cache ngay trên chip. Key-value cache có thể xem như bộ nhớ ngắn hạn của mô hình trong quá trình hội thoại, lưu các biểu diễn trung gian để không phải tính lại từ đầu cho mỗi token mới. Ngoài ra, TPU 8i loại bỏ SparseCore và thay bằng CAE, viết tắt của collective acceleration engine. Đây là khối tăng tốc chuyên xử lý giao tiếp tập thể giữa nhiều chip, như all-reduce hoặc all-gather, hai thao tác đồng bộ dữ liệu rất phổ biến trong AI phân tán. Khi đẩy phần giao tiếp này sang phần cứng chuyên dụng, các tensor core, tức lõi tính toán chính cho AI, sẽ bớt phải chờ dữ liệu và làm việc hiệu quả hơn.

Nhu cầu tối ưu giao tiếp tập thể đặc biệt tăng mạnh khi các mô hình MoE, hay mixture-of-experts, ngày càng phổ biến. Đây là kiến trúc chia mô hình thành nhiều “chuyên gia” nhỏ; với mỗi token, chỉ một số chuyên gia được kích hoạt thay vì toàn bộ mô hình. Lợi ích là mô hình có thể mở rộng quy mô mà không nhất thiết phải tăng tương ứng nhu cầu bộ nhớ và tính toán cho mọi lượt suy luận. Nhưng cái giá phải trả là lưu lượng giao tiếp giữa các chip trở nên khó đoán hơn, vì mỗi token có thể cần gọi đến các chuyên gia nằm trên những bộ tăng tốc khác nhau. Google nói CAE giúp giảm độ trễ của các tác vụ giao tiếp tập thể tới 5 lần, từ đó cải thiện hiệu quả kinh tế khi một cụm máy có thể phục vụ nhiều người dùng hơn.

Bên cạnh chip TPU 8i, Google còn giới thiệu topology mạng mới tên Boardfly cho các cụm suy luận. Topology, hay cấu trúc liên kết mạng, là cách các chip và thiết bị mạng được nối với nhau. Google cho biết Boardfly gợi nhớ đến Dragonfly, một kiểu thiết kế quen thuộc trong các hệ thống HPC, và cho phép kết nối 1.152 chip, trong đó 1.024 chip hoạt động tại một thời điểm, thông qua chuyển mạch quang. Lợi ích chính là giảm số bước nhảy tối đa giữa hai chip từ 16 xuống còn 7 so với mạng 3D torus truyền thống. 3D torus là kiểu liên kết lưới ba chiều, hiệu quả ở một số bài toán nhưng có thể gây độ trễ cao hơn khi lưu lượng giao tiếp thay đổi thất thường như trong MoE hoặc các mô hình suy luận nhiều bước.

Thông điệp lớn nhất từ lần ra mắt này là cuộc chiến AI không còn xoay quanh một con chip mạnh nhất, mà là ai xây được hạ tầng hoàn chỉnh tốt hơn. Google đang kết hợp chip TPU chuyên biệt, CPU Arm Axion, mạng quang, lưu trữ tốc độ cao và nền tảng AI Hypercomputer để tạo ra một hệ sinh thái phục vụ huấn luyện lẫn triển khai LLM ở quy mô lớn. AI Hypercomputer có thể hiểu là gói hạ tầng trọn bộ gồm tính toán, mạng, lưu trữ và phần mềm để khách hàng triển khai AI mà không phải tự lắp ghép từng thành phần. Cả TPU 8t và TPU 8i dự kiến sẽ có mặt rộng rãi trên Google Cloud vào cuối năm nay, cho thấy Google muốn biến lợi thế hạ tầng nội bộ thành vũ khí thương mại trong cuộc đua với Nvidia, AWS và các đối thủ đám mây khác.