TPU (Tensor Processing Unit)

Định nghĩa

TPU (viết tắt của Tensor Processing Unit) là một loại đơn vị xử lý chuyên biệt được thiết kế từ đầu để thực hiện các phép toán trên tensor — cấu trúc dữ liệu đa chiều cơ bản trong học máy và học sâu. Khác với các bộ vi xử lý chung như CPU hay GPU vốn có kiến trúc đa năng nhằm đáp ứng nhiều loại tác vụ khác nhau, TPU được tối ưu hóa sâu sắc cho các phép tính ma trận và vector quy mô lớn, đặc biệt là nhân ma trận (matrix multiplication), tích vô hướng (dot product), và các hàm kích hoạt phi tuyến phổ biến như ReLU, sigmoid hay softmax. Thuật ngữ "tensor" trong tên gọi không chỉ mang ý nghĩa toán học thuần túy — chỉ đối tượng tổng quát hóa của vô hướng, vectơ và ma trận — mà còn phản ánh cách thức biểu diễn dữ liệu và tham số trong các khung làm việc học sâu như TensorFlow, nơi mọi biến đều được tổ chức dưới dạng tensor với số chiều tùy ý.

Từ góc độ kiến trúc hệ thống, TPU không phải là một bộ vi xử lý độc lập theo nghĩa truyền thống, mà là một khối xử lý đồng tiến (accelerator) được tích hợp vào hệ sinh thái phần cứng và phần mềm của Google, hoạt động như một đơn vị mở rộng cho hệ thống máy chủ hoặc trung tâm dữ liệu. Nó không chạy hệ điều hành riêng, cũng không hỗ trợ lập trình cấp thấp theo kiểu C/C++ trực tiếp; thay vào đó, nó được điều khiển thông qua các trình biên dịch và runtime đặc biệt như XLA (Accelerated Linear Algebra), chuyển đổi các biểu đồ tính toán (computation graphs) từ ngôn ngữ bậc cao thành chuỗi lệnh vi mã (microcode) phù hợp với logic mạch cứng bên trong. Việc đặt tên là "Unit" thay vì "Processor" hay "Chip" cũng thể hiện rõ định hướng thiết kế: TPU là một thành phần chức năng trong một hệ thống lớn hơn, không tồn tại độc lập mà luôn gắn kết chặt chẽ với phần mềm điều khiển và hạ tầng lưu trữ/di chuyển dữ liệu.

Một điểm then chốt trong định nghĩa TPU là sự phân biệt rõ ràng giữa vai trò chuyên biệt của nó so với các công nghệ tăng tốc khác. Trong khi GPU ban đầu được thiết kế cho đồ họa nhưng sau đó được khai thác cho tính toán chung (GPGPU), thì TPU là sản phẩm đầu tiên trên thế giới được xây dựng từ con số không với mục tiêu duy nhất: phục vụ học máy. Điều này dẫn đến những lựa chọn kiến trúc căn bản khác biệt — chẳng hạn như việc ưu tiên băng thông bộ nhớ cao và độ trễ thấp hơn là khả năng linh hoạt trong lập trình, hay việc sử dụng số học cố định (fixed-point arithmetic) thay vì dấu phẩy động (floating-point) ở các thế hệ đầu nhằm tiết kiệm diện tích chip và năng lượng.

Lịch sử và nguồn gốc

TPU được Google công bố lần đầu tiên vào tháng 5 năm 2016 tại hội nghị I/O thường niên, đánh dấu một bước ngoặt quan trọng trong lịch sử phần cứng AI. Dự án phát triển TPU bắt đầu từ năm 2013 trong nội bộ Google Brain — nhóm nghiên cứu trí tuệ nhân tạo do Geoffrey Hinton và Andrew Ng từng cộng tác — khi các kỹ sư nhận ra rằng việc huấn luyện các mô hình mạng nơ-ron sâu như Inception hay GNMT (Google Neural Machine Translation) trên GPU NVIDIA đang gặp giới hạn nghiêm trọng về hiệu suất và hiệu quả năng lượng. Các mô hình ngày càng lớn, dữ liệu huấn luyện ngày càng phong phú, trong khi tốc độ xử lý và mức tiêu thụ điện năng của GPU không còn đáp ứng đủ nhu cầu mở rộng quy mô hệ thống dịch máy tự động, nhận dạng giọng nói và tìm kiếm hình ảnh của Google.

Quá trình phát triển được dẫn dắt bởi nhóm kỹ sư phần cứng và phần mềm do Norman Jouppi — nhà khoa học nghiên cứu cấp cao tại Google — đứng đầu. Đội ngũ đã tiến hành phân tích sâu các biểu đồ tính toán thực tế từ hàng chục nghìn tác vụ học máy đang chạy trên hạ tầng Google Cloud và Data Centers, từ đó xác định các mẫu truy cập bộ nhớ, tỷ lệ phép toán nhân-ma-trận so với các phép toán điều khiển, và đặc điểm phân bố độ chính xác yêu cầu. Kết quả phân tích cho thấy hơn 90% thời gian tính toán tập trung vào các phép nhân và cộng ma trận, trong khi các tác vụ điều khiển luồng (branching), quản lý bộ nhớ động hay xử lý ngoại lệ gần như không xuất hiện. Đây là cơ sở để thiết kế một kiến trúc không cần đơn vị logic số học (ALU) mạnh mẽ, không cần bộ nhớ cache phức tạp, mà thay vào đó tập trung vào một mảng nhân ma trận khổng lồ (matrix multiply unit) cùng với một hệ thống bộ nhớ đệm (on-chip memory) được sắp xếp theo cấu trúc systolic array — một mô hình xử lý song song đồng bộ từng bước, nơi dữ liệu chảy liên tục giữa các ô xử lý theo một chu kỳ cố định.

Các mốc quan trọng trong lịch sử TPU bao gồm: phiên bản TPU v1 ra mắt năm 2015 được triển khai nội bộ để hỗ trợ dịch máy tự động Google Translate và hệ thống đề xuất YouTube; TPU v2 (2017) với khả năng huấn luyện mô hình, tích hợp công nghệ interconnect 2D torus cho cụm nhiều chip; TPU v3 (2018) nâng cấp công suất làm mát và băng thông bộ nhớ, đạt hiệu năng hơn 400 teraFLOPS (FP16); TPU v4 (2021) tích hợp hỗ trợ số học bfloat16, tăng cường bảo mật phần cứng và khả năng mở rộng lên hàng nghìn chip trong một hệ thống; và mới đây nhất là TPU v5e (2023) và TPU v5p (2024), được tối ưu cho cả suy luận thời gian thực và huấn luyện mô hình ngôn ngữ lớn (LLM) với hiệu suất năng lượng vượt trội. Ngoài ra, Google còn phát triển TPU dành riêng cho thiết bị biên (Edge TPU) từ năm 2018, nhằm triển khai mô hình học sâu trên thiết bị IoT, camera thông minh hay thiết bị y tế cầm tay.

Đặc điểm và tính chất

TPU là một hệ thống vi mạch phức tạp, kết hợp giữa kiến trúc phần cứng đặc thù và phần mềm điều khiển tích hợp sâu. Đặc điểm nổi bật nhất là thiết kế theo nguyên tắc domain-specific architecture (DSA), tức là kiến trúc chuyên biệt cho một miền ứng dụng cụ thể — trong trường hợp này là học máy. Điều này thể hiện rõ qua việc loại bỏ hoặc giảm thiểu các thành phần không cần thiết trong kiến trúc chung như đơn vị điều khiển luồng phức tạp, bộ nhớ cache đa cấp, hay bộ quản lý bộ nhớ ảo. Thay vào đó, TPU tập trung tài nguyên vào các khối chức năng then chốt: khối nhân ma trận (Matrix Multiply Unit), khối tích lũy (Accumulator), khối bộ nhớ đệm tích hợp (Unified Buffer), và khối điều khiển vi mã (Microcode Controller).

• Kiến trúc systolic array: Toàn bộ khối nhân ma trận được xây dựng dưới dạng mảng systolic — một cấu trúc xử lý song song nơi mỗi ô (cell) thực hiện một phép nhân và cộng cục bộ, sau đó truyền kết quả sang ô lân cận theo chu kỳ đồng hồ cố định. Kiến trúc này cho phép dữ liệu di chuyển một cách có trật tự và dự đoán được, giảm đáng kể chi phí vận chuyển dữ liệu — yếu tố chiếm tới 90% năng lượng tiêu thụ trong các hệ thống xử lý chung.
• Bộ nhớ tích hợp cao băng thông: TPU sử dụng bộ nhớ SRAM tích hợp trực tiếp trên chip (on-die memory), với dung lượng từ vài chục MB đến hơn 100 MB tùy thế hệ. Bộ nhớ này được chia thành các bank độc lập, cho phép truy cập đồng thời nhiều luồng dữ liệu, đảm bảo khối nhân ma trận luôn được cung cấp dữ liệu đầy đủ mà không bị nghẽn (data starvation). Băng thông bộ nhớ có thể đạt tới hơn 1 TB/s ở các thế hệ mới — cao gấp nhiều lần so với băng thông DRAM ngoài chip của GPU tương đương.
• Hệ thống số học tối ưu: Các thế hệ TPU đầu tiên sử dụng số học INT8 (số nguyên 8 bit) cho suy luận và bfloat16 (số dấu phẩy động 16 bit với độ chính xác thấp hơn FP32 nhưng ổn định hơn FP16) cho huấn luyện. Việc từ bỏ IEEE 754 FP32 giúp giảm diện tích chip, tăng mật độ tính toán và giảm tiêu thụ điện năng. Các khối số học được thiết kế cứng (hardwired), không có khả năng lập trình lại như ALU trên CPU, nhưng đạt hiệu suất cực cao cho các phép toán mục tiêu.
• Giao diện kết nối chuyên biệt: TPU không sử dụng chuẩn PCIe như GPU, mà kết nối qua các giao diện tùy chỉnh như Inter-Chip Interconnect (ICI) hoặc Optical ICI (trong TPU v4), cho phép tốc độ truyền dữ liệu giữa các chip lên tới hàng chục TB/s. Điều này cho phép xây dựng các cụm TPU quy mô lớn (hàng trăm hoặc hàng nghìn chip) hoạt động như một siêu máy tính thống nhất.

Một đặc điểm kỹ thuật quan trọng khác là khả năng tích hợp phần mềm-hardware. Mỗi TPU đi kèm với một trình biên dịch XLA được tối ưu hóa, có khả năng phân tích biểu đồ tính toán và tái sắp xếp thứ tự thực thi nhằm tận dụng tối đa băng thông bộ nhớ và độ song song của mảng systolic. Ngoài ra, hệ thống runtime (như TensorFlow Serving hoặc JAX) cung cấp lớp trừu tượng hóa cao, che giấu hoàn toàn độ phức tạp của phần cứng, cho phép nhà nghiên cứu chỉ cần viết mã Python mà không cần hiểu chi tiết về vi kiến trúc.

Phân loại

TPU cho trung tâm dữ liệu (Cloud TPU)

Đây là dòng TPU chính, được thiết kế để triển khai trong các trung tâm dữ liệu của Google Cloud. Bao gồm các thế hệ v1 đến v5p, mỗi thế hệ được cải tiến về hiệu năng, hiệu quả năng lượng, khả năng mở rộng và hỗ trợ phần mềm. TPU v1 và v2 chủ yếu dùng cho suy luận và huấn luyện mô hình vừa và nhỏ; TPU v3 và v4 mở rộng khả năng xử lý mô hình ngôn ngữ lớn với bộ nhớ tích hợp lớn hơn và hỗ trợ bfloat16; TPU v5e (enhanced) tập trung vào hiệu năng trên watt cho suy luận thời gian thực, trong khi TPU v5p (premier) là phiên bản cao cấp nhất, hỗ trợ cả huấn luyện LLM quy mô hàng chục tỷ tham số với độ trễ cực thấp.

TPU cho thiết bị biên (Edge TPU)

Edge TPU là phiên bản thu nhỏ, tiêu thụ điện năng thấp (dưới 2W), được tích hợp vào các module phần cứng như Coral USB Accelerator hoặc Dev Board. Nó được thiết kế để chạy các mô hình học sâu đã được lượng tử hóa (quantized) và cắt tỉa (pruned), chủ yếu ở định dạng TensorFlow Lite. Edge TPU không hỗ trợ huấn luyện, chỉ thực hiện suy luận, nhưng có khả năng hoạt động độc lập trên thiết bị không có kết nối mạng, đáp ứng yêu cầu về quyền riêng tư, độ trễ và độ tin cậy trong các ứng dụng như giám sát an ninh, chẩn đoán y tế tại chỗ hay điều khiển robot công nghiệp.

TPU dành riêng cho nghiên cứu (Research TPU)

Google cũng cung cấp chương trình TPU Research Cloud (TRC), cấp quyền truy cập miễn phí vào các cụm TPU cho các nhà nghiên cứu học thuật. Các cụm này thường sử dụng TPU v3 hoặc v4 với cấu hình đặc biệt, được tối ưu hóa cho các thí nghiệm khoa học như mô phỏng vật lý lượng tử, phân tích gen hoặc mô hình hóa khí hậu — những lĩnh vực đòi hỏi tính toán tensor quy mô lớn nhưng không thuộc phạm vi thương mại truyền thống.

Cơ chế hoạt động

Cơ chế hoạt động của TPU dựa trên nguyên lý xử lý dữ liệu theo luồng (streaming processing) trong môi trường được kiểm soát chặt chẽ bởi phần mềm biên dịch. Khi một mô hình được gửi đến TPU, trình biên dịch XLA phân tích biểu đồ tính toán, xác định thứ tự thực thi, phân bổ tensor vào bộ nhớ tích hợp, và tạo ra chuỗi lệnh vi mã (microcode) điều khiển từng ô trong mảng systolic. Dữ liệu đầu vào (ví dụ: ảnh, văn bản mã hóa) được tải vào Unified Buffer, sau đó được phân phối đồng thời đến các hàng và cột của mảng nhân ma trận theo sơ đồ di chuyển dữ liệu cố định. Mỗi ô thực hiện một phép nhân và cộng cục bộ, kết quả được tích lũy trong thanh ghi tích lũy cục bộ, rồi truyền tiếp sang ô kế tiếp. Sau một số chu kỳ đồng hồ nhất định, kết quả cuối cùng được đọc ra và đưa vào khối xử lý hàm kích hoạt hoặc khối chuẩn hóa lớp (layer normalization). Toàn bộ quá trình diễn ra mà không cần can thiệp của bộ điều khiển trung tâm — mọi thứ được điều phối bởi tín hiệu đồng bộ hóa và cấu trúc mạch cứng.

Ứng dụng thực tế

TPU được sử dụng rộng rãi trong các dịch vụ sản phẩm của Google như Google Search (tối ưu hóa kết quả tìm kiếm bằng mô hình xếp hạng học sâu), Gmail (phát hiện thư rác và gợi ý trả lời tự động), Google Photos (nhận dạng đối tượng và phân cụm ảnh), và Google Assistant (nhận dạng giọng nói và tổng hợp tiếng nói). Trong lĩnh vực khoa học, TPU hỗ trợ Dự án AlphaFold của DeepMind trong việc dự đoán cấu trúc protein — một bài toán yêu cầu hàng triệu phép tính tensor mỗi giây. Ngoài ra, các viện nghiên cứu như CERN hay NASA cũng sử dụng TPU để phân tích dữ liệu va chạm hạt hoặc mô phỏng khí quyển sao Hỏa. Trong y học, các bệnh viện đối tác với Google Health sử dụng TPU để huấn luyện mô hình phát hiện ung thư vú trên ảnh chụp X-quang với độ chính xác vượt trội so với bác sĩ chuyên khoa.

Ưu điểm và hạn chế

Ưu điểm nổi bật nhất của TPU là hiệu suất năng lượng — đạt tới hơn 100 teraFLOPS/W, cao gấp 2–3 lần so với GPU hàng đầu cùng thời. Nhờ kiến trúc chuyên biệt, TPU cũng có độ trễ thấp hơn đáng kể trong các tác vụ suy luận thời gian thực, và khả năng mở rộng ngang (horizontal scaling) vượt trội khi kết nối hàng nghìn chip. Về mặt phần mềm, hệ sinh thái tích hợp với TensorFlow và JAX cung cấp trải nghiệm phát triển liền mạch, giảm thiểu rào cản kỹ thuật cho nhà nghiên cứu.

Hạn chế chính là tính chuyên biệt: TPU không thể chạy các ứng dụng chung như hệ điều hành, trình duyệt hay phần mềm văn phòng; nó chỉ hoạt động hiệu quả với các mô hình được viết bằng các framework hỗ trợ (chủ yếu là TensorFlow và JAX), và yêu cầu người dùng phải hiểu rõ cách tối ưu hóa biểu đồ tính toán. Ngoài ra, chi phí đầu tư ban đầu cho hạ tầng TPU rất cao, và khả năng tương thích ngược giữa các thế hệ không hoàn toàn đảm bảo. Việc debug và profiling cũng khó khăn hơn so với CPU/GPU do thiếu công cụ phân tích cấp phần cứng phổ biến.

Lưu ý quan trọng

Khi triển khai TPU, người dùng cần lưu ý rằng hiệu năng thực tế phụ thuộc rất lớn vào cách thiết kế mô hình và tối ưu hóa biểu đồ tính toán — một mô hình không được thiết kế phù hợp có thể chỉ khai thác được dưới 30% hiệu năng lý thuyết. Việc lượng tử hóa sai hoặc cắt tỉa quá mức sẽ làm suy giảm độ chính xác nghiêm trọng. Ngoài ra, do TPU không hỗ trợ quản lý bộ nhớ ảo, nên kích thước tensor đầu vào phải vừa khít với dung lượng bộ nhớ tích hợp; nếu vượt quá, hệ thống sẽ báo lỗi OOM (out-of-memory) thay vì hoán chuyển sang bộ nhớ ngoài như GPU. Cuối cùng, cần hiểu rõ giới hạn phần cứng: TPU không hỗ trợ các phép toán không tuyến tính phức tạp như FFT hay giải phương trình vi phân, và không thể thay thế hoàn toàn CPU trong các hệ thống tổng hợp.