TPU là gì? Khác CPU và GPU thế nào, khi nào nên dùng

TPU (Tensor Processing Unit) là con chip ASIC do Google tự phát triển, chuyên để tăng tốc các phép tính ma trận trong học máy. Khác với CPU làm được nhiều việc hoặc GPU làm tốt việc song song nói chung, TPU được thiết kế riêng cho tính toán mạng nơ-ron. Bài này giải thích TPU hoạt động ra sao, khác CPU và GPU ở đâu, khi nào hợp dùng và nếu tự dựng máy để học AI tại nhà thì nên chọn gì.

TPU là gì?

Theo Google Cloud, TPU là vi mạch tích hợp chuyên dụng (ASIC) do Google phát triển để tăng tốc các tác vụ học máy. Chip được thiết kế cho các phép toán ma trận lớn thường gặp trong thuật toán học máy và có bộ nhớ băng thông cao (HBM) ngay trên chip, giúp chạy mô hình và kích thước batch lớn hơn. Nhiều TPU có thể nối thành nhóm gọi là slice để mở rộng tải. Nếu bạn chưa quen khái niệm học máy, bài Machine Learning là gì giải thích phần nền tảng trước.

TPU hoạt động như thế nào?

Mạng nơ-ron cần cực nhiều phép nhân và cộng giữa dữ liệu đầu vào và tham số, được tổ chức thành phép nhân ma trận. Google Cloud giải thích CPU là bộ xử lý đa dụng theo kiến trúc von Neumann: mỗi phép tính phải đọc ghi kết quả qua bộ nhớ nên bị thắt cổ chai. GPU dùng nhiều đơn vị ALU chạy song song (Google nêu khoảng 2.500 đến 5.000 ALU mỗi bộ xử lý) nên thông lượng cao hơn CPU nhiều, nhưng vẫn là bộ xử lý đa dụng nên còn truy cập thanh ghi hoặc bộ nhớ chia sẻ cho mỗi phép tính.

TPU đi theo hướng khác: vì biết trước các bước xử lý ma trận, Google nối hàng nghìn bộ nhân và cộng trực tiếp với nhau thành một ma trận vật lý gọi là systolic array. Ở Cloud TPU v2, mỗi bộ xử lý có hai systolic array 128 x 128, tổng 32.768 ALU cho số thực 16 bit. Kết quả nhân được chuyển thẳng sang bộ nhân kế tiếp để cộng dồn, nên trong quá trình này không cần truy cập bộ nhớ. Đổi lại, TPU không chạy được các việc đa dụng như CPU.

Khi nào TPU phù hợp, khi nào không?

Trang tài liệu của Google Cloud nêu rõ các trường hợp TPU hợp với những mô hình có phép tính ma trận chiếm chủ đạo, vòng huấn luyện chính không có phép toán PyTorch hoặc JAX tùy biến, huấn luyện kéo dài nhiều tuần đến nhiều tháng, mô hình lớn có batch hiệu dụng lớn, hoặc mô hình có embedding cực lớn như xếp hạng và gợi ý. Ngược lại, TPU không phù hợp với chương trình đại số tuyến tính có nhiều rẽ nhánh hoặc phép tính theo từng phần tử, tác vụ cần số học độ chính xác cao, mạng có phép toán tùy biến trong vòng huấn luyện chính, hay mô hình có hình dạng tensor thay đổi vì đồ thị được biên dịch theo batch đầu tiên.

Điểm cần nhớ là TPU là dịch vụ bạn thuê qua đám mây (Google nêu các đường truy cập qua Compute Engine, Google Kubernetes Engine), không phải linh kiện bán lẻ để lắp vào máy. Vì thế người học AI tại nhà gần như chỉ có đường chọn GPU.

Từ trái sang: PC AI Ryzen 9 9950X với RTX PRO 5000 Blackwell 48GB, Agent Computer 9950X3D với RTX 5080, Agent Computer 9900X3D với RTX 5080

Muốn học và chạy AI tại nhà thì chọn gì?

Nếu bạn muốn một máy chạy được mô hình mà không phải thuê đám mây, nhóm bộ PC chạy AI của Nguyễn Công đã cân đối sẵn CPU, RAM và card. Nhu cầu chạy tác tử AI có thể xem Agent Computer, còn người học di động có laptop AI. Tự ghép máy thì bắt đầu từ card NVIDIA RTX 5000 Series, và khi cần nhiều VRAM hơn cho mô hình lớn thì xem RTX PRO 6000 Blackwell. Đừng quên RAM DDR5 và ổ SSD đủ dung lượng, vì bộ dữ liệu lớn sẽ chiếm bộ nhớ trước khi đến GPU. Thông số từng sản phẩm bạn nên đối chiếu trên trang bán và trang của hãng.

Từ trái sang: NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition, Gigabyte GeForce RTX 5060 Windforce Max OC 8G, Asus Dual GeForce RTX 5060 Ti 16GB GDDR7 OC Edition

Câu hỏi thường gặp

TPU là gì?

TPU (Tensor Processing Unit) là vi mạch ASIC do Google phát triển để tăng tốc các phép tính ma trận trong học máy.

TPU viết tắt của gì?

Tensor Processing Unit, nghĩa là đơn vị xử lý tensor.

TPU khác GPU thế nào?

GPU là bộ xử lý đa dụng chạy song song nhiều ALU, dùng được cho đồ họa lẫn AI. TPU là chip chuyên cho mạng nơ-ron, có systolic array nối trực tiếp bộ nhân và cộng, nhưng không chạy được tác vụ đa dụng.

TPU khác CPU thế nào?

CPU đa dụng và linh hoạt nhưng mỗi phép tính phải qua bộ nhớ. TPU đặc thù cho ma trận nên bỏ được bước truy cập bộ nhớ trong lúc nhân ma trận.

Systolic array là gì?

Là ma trận vật lý gồm hàng nghìn bộ nhân và cộng nối trực tiếp với nhau, kết quả chuyển thẳng sang bộ kế tiếp mà không cần truy cập bộ nhớ.

Tôi có mua TPU để lắp vào PC được không?

Theo mô tả của Google, TPU được cung cấp dưới dạng dịch vụ đám mây qua Compute Engine, Google Kubernetes Engine. Bạn không lắp TPU của Google vào PC như một card đồ họa.

Người mới học AI nên dùng TPU hay GPU?

Với máy cá nhân, GPU NVIDIA là lựa chọn thực tế vì bạn mua được, lắp được và có nhiều tài liệu hướng dẫn.

TPU hợp với loại mô hình nào?

Mô hình có phép tính ma trận chiếm chủ đạo, huấn luyện dài ngày, batch lớn, hoặc có embedding rất lớn như xếp hạng và gợi ý.

TPU không hợp với trường hợp nào?

Chương trình nhiều rẽ nhánh hoặc phép tính từng phần tử, cần độ chính xác số học cao, có phép toán tùy biến trong vòng huấn luyện, hoặc có hình dạng tensor thay đổi.

TPU có tiết kiệm điện hơn không?

Google mô tả TPU xử lý phép nhân ma trận mạng nơ-ron nhanh hơn trong khi tốn ít điện và chiếm ít diện tích hơn, đây là mô tả của chính Google, không phải kết quả bạn tự đo.

Học máy trên máy tính cá nhân cần cấu hình gì?

Một card NVIDIA đủ VRAM, CPU đủ lõi, RAM và SSD đủ dung lượng cho bộ dữ liệu. Bạn có thể xem các bộ PC chạy AI hoặc card RTX 5000 Series để chọn.

Nên mua bộ PC chạy AI lắp sẵn hay tự ghép?

Lắp sẵn tiết kiệm thời gian cân đối CPU, RAM, card. Tự ghép linh hoạt hơn nếu bạn đã rõ mô hình và dung lượng VRAM cần dùng.