- TPU là chip tùy biến do Google phát triển dành cho huấn luyện và suy luận AI quy mô lớn, với triết lý thiết kế khác với GPU
- Nhấn mạnh khả năng mở rộng và hiệu quả năng lượng, đồng thời cùng thiết kế phần cứng (ví dụ: cấu hình system-on-chip, bộ nhớ on-chip lớn) và phần mềm (trình biên dịch XLA)
- Cấu trúc cốt lõi là mảng systolic, pipeline và phương thức biên dịch trước, được tối ưu cho phần lớn các phép toán deep learning (đặc biệt là nhân ma trận)
- Nhờ công nghệ OCI và OCS, có thể cấu hình node linh hoạt, xử lý song song hiệu năng cao và lựa chọn nhiều topology khác nhau
- Hệ thống TPU có thể mở rộng theo tầng từ một chip đơn lẻ đến multi-pod siêu lớn, cho phép huấn luyện các mô hình AI quy mô lớn và sử dụng tài nguyên hiệu quả
Tổng quan và bối cảnh
- TPU là chip dựa trên ASIC do Google phát triển để đáp ứng việc mở rộng quy mô lớn của các dịch vụ AI
- Ban đầu Google cân nhắc giữa GPU, FPGA và ASIC, nhưng từ năm 2013, khi việc ứng dụng deep learning như tìm kiếm bằng giọng nói mở rộng, nhu cầu về phần cứng tùy biến ngày càng tăng
- Hiện nay TPU đóng vai trò thiết yếu trong các dịch vụ AI chủ lực như Gemini, Veo, đồng thời được sử dụng rộng rãi cho các mô hình gợi ý (DLRM)
Kiến trúc chip TPU đơn lẻ
Cấu hình cơ bản
- Theo chuẩn TPUv4, một chip đơn có 2 TensorCore (TPU chuyên cho suy luận có 1)
- Mỗi TensorCore được kết nối với bộ nhớ CMEM (128MiB) và HBM (32GiB)
Cấu trúc bên trong TensorCore
- Matrix Multiply Unit (MXU): dựa trên mảng systolic 128x128, chuyên trách nhân ma trận
- Vector Unit (VPU): thực hiện các phép toán phần tử thông thường
- Vector Memory (VMEM; 32MiB): sao chép dữ liệu từ HBM để dùng làm bộ nhớ chuẩn bị công việc
- Scalar Unit + Scalar Memory (SMEM; 10MiB): phụ trách luồng điều khiển, phép toán scalar và quản lý địa chỉ bộ nhớ
Khác biệt về cấu trúc so với GPU
- TPU có bộ nhớ on-chip (CMEM, VMEM, SMEM) lớn hơn nhiều so với GPU
- Dung lượng HBM thì GPU lớn hơn, và số lượng lõi tính toán của GPU cũng nhiều hơn đáng kể
- Theo chuẩn TPUv5p, hiệu năng đạt 500 TFLOPs/sec mỗi chip, và toàn bộ pod (8960 chip) đạt khoảng 4.45 ExaFLOPs/sec
Triết lý thiết kế TPU
1. Mảng systolic và pipeline
- Mảng systolic là cấu trúc trong đó các phần tử xử lý (PE) được sắp xếp thành mảng và truyền kết quả tính toán sang phần tử lân cận
- Sau khi dữ liệu được đưa vào, có thể tính toán liên tục mà không cần điều khiển bổ sung; việc đọc/ghi bộ nhớ chỉ xảy ra ở đầu vào và đầu ra
- Chủ yếu được tối ưu cho nhân ma trận, convolution
- Thông qua xử lý pipeline, việc tính toán và di chuyển dữ liệu diễn ra đồng thời để tối ưu thông lượng
Nhược điểm của mảng systolic - độ thưa (sparsity)
- Tất cả phần tử tính toán luôn hoạt động nên không phù hợp với ma trận thưa
- Nếu các mô hình DL trong tương lai theo đuổi tính thưa bất quy tắc, đây có thể trở thành một giới hạn
2. Biên dịch trước (AoT) và giảm thiểu cache
- Nhờ đồng thiết kế TPU-XLA, sự phụ thuộc vào cache vốn cần cho truy cập bộ nhớ bất quy tắc được giảm thiểu, từ đó tiết kiệm năng lượng
- Trình biên dịch XLA phân tích đồ thị tính toán để xác định trước mẫu truy cập bộ nhớ, vận hành chủ yếu dựa trên scratchpad memory thay vì cache
@jit của JAX là dạng trung gian giữa JIT và AoT: khi chạy lần đầu sẽ tạo đồ thị tĩnh rồi tiến hành biên dịch AoT bằng XLA- Khi hình dạng đầu vào thay đổi thì cần biên dịch lại, và vẫn kém hiệu quả với padding/loop động
- Đạt hiệu quả năng lượng cao nhưng nhược điểm là thiếu linh hoạt
Hiệu quả năng lượng của TPUv4
- Các chip hiện đại sử dụng bộ nhớ HBM3 để giảm tiêu thụ năng lượng
- Vì năng lượng cho thao tác bộ nhớ có thể cao hơn hàng chục đến hàng trăm lần so với tính toán, việc giảm truy cập bộ nhớ có thể cải thiện hiệu suất rất lớn
Kiến trúc TPU đa chip
Cấp độ tray (Tray/Board; 4 chip)
- 1 tray gồm 4 chip TPU (8 TensorCore) và CPU Host
- Kết nối Host↔Chip là PCIe, còn Chip↔Chip là Inter-Core Interconnect (ICI), cung cấp băng thông cao hơn
Cấp độ rack (Rack; 4x4x4=64 chip)
- 1 rack gồm 64 chip TPU, được kết nối bằng ICI và OCS (Optical Circuit Switching) theo 3D torus (4x4x4)
- Google phân biệt rack, pod và slice
- Rack: đơn vị vật lý 64 chip (= khối lập phương)
- Pod: đơn vị tối đa có thể kết nối bằng ICI và OCS (ví dụ: TPUv4 = 4096 chip = 64 rack)
- Slice: đơn vị trừu tượng có thể cấu hình tùy ý từ 4 chip đến Superpod
Lợi ích của OCS
- Wraparound: biến mỗi trục thành vòng (1D torus), giúp giảm số hop di chuyển tệ nhất giữa các node
- Slice không liền kề linh hoạt: nhờ cấu trúc chuyển mạch OCS, các node nằm tách rời về vật lý vẫn có thể được cấu hình thành một slice, thuận tiện cho việc sử dụng tài nguyên và bảo trì
- Topology xoắn: với cùng số lượng chip (cố định x, y, z), có thể thay đổi cấu trúc kết nối để tối ưu tốc độ cho các mẫu tính toán cụ thể (ví dụ: twisted torus)
Ví dụ sử dụng topology
- Dạng khối lập phương: phù hợp cho song song dữ liệu/tensor (băng thông tối đa)
- Dạng đường thẳng (kiểu xì gà): có lợi cho song song pipeline
- Twisted torus: cải thiện tốc độ khi cần giao tiếp all-to-all (ví dụ: tensor parallel)
Superpod (Full Pod/Superpod; TPUv4: 4096 chip/64 rack)
- Kết nối nhiều rack bằng ICI và OCS để xây dựng hệ thống siêu lớn
- Tùy theo loại topology của slice, băng thông giao tiếp và hiệu năng xử lý song song sẽ khác nhau
- Nhờ OCS, có được tính linh hoạt như slice không liền kề hay topology xoắn
Cấp độ multi-pod (Multi-pod/Multislice; TPUv4: hơn 4096 chip)
- Nhiều pod được kết nối qua mạng trung tâm dữ liệu (DCN) để cung cấp hạ tầng huấn luyện quy mô lớn, nhưng băng thông thấp hơn ICI
- Việc huấn luyện PaLM sử dụng 2 pod (6144 TPUv4), với tổng cộng tài nguyên của 6 pod được quản lý
- Khi phát triển mô hình quy mô lớn, trình biên dịch XLA sẽ tối ưu mẫu giao tiếp tùy theo các giá trị cấu hình của nhà nghiên cứu (như chiều song song)
- XLA chèn các phép toán giao tiếp giữa từng slice và pod, giúp triển khai huấn luyện phân tán quy mô lớn với thay đổi mã tối thiểu
Ví dụ về phần cứng thực tế và sơ đồ
- Rack TPU: một đơn vị gồm 3D torus 4x4x4, mỗi hàng có 2 tray (8 chip)
- Tray TPUv4: thực tế có 4 cổng PCIe (mỗi TPU một cổng)
- Chip TPUv4: ASIC ở trung tâm, xung quanh là 4 stack HBM (theo chuẩn 2 TensorCore)
- Sơ đồ mặt bằng chip TPUv4i (dạng suy luận): 1 TensorCore, diện tích CMEM chiếm phần lớn
Kết
- Bày tỏ sự hỗ trợ nghiên cứu từ Google TPU Research Cloud (TRC)
Tài liệu tham khảo
Chưa có bình luận nào.