Google phát hành Hướng dẫn đào tạo nhà phát triển TPU Ironwood thế hệ thứ 7, chi tiết tối ưu hóa hiệu suất hệ thống

Tin tức ME News, ngày 2 tháng 4 (UTC+8), Google chính thức gần đây đã công bố hướng dẫn đào tạo dành cho nhà phát triển cho thế hệ thứ bảy Ironwood TPU. Tài liệu này nhằm giúp các nhà phát triển tận dụng tối đa hiệu năng ở cấp hệ thống của Ironwood TPU để đào tạo và triển khai hiệu quả các mô hình AI tiên tiến. Ironwood TPU là cơ sở hạ tầng AI tùy chỉnh được thiết kế để đáp ứng nhu cầu tính toán cho các mô hình tham số hàng nghìn tỉ; thông qua các công nghệ như liên kết liên chip (ICI), bộ chuyển mạch đường quang (OCS), mạng trung tâm dữ liệu (DCN) và bộ nhớ băng thông cao được tổng hợp (HBM), nó xây dựng nên một hệ thống hoàn chỉnh hỗ trợ lên đến 9.216 chip. Bài viết mô tả chi tiết nhiều chiến lược tối ưu then chốt cho phần cứng này, bao gồm: sử dụng đơn vị nhân ma trận (MXU) hỗ trợ nguyên sinh huấn luyện FP8 để nâng cao thông lượng; sử dụng bộ thư viện kernel JAX được tối ưu riêng cho TPU là Tokamax, thông qua “attention bắn tung” và “Megablox phân nhóm phép nhân ma trận” để xử lý các tensor không đều trong mô hình ngữ cảnh dài và mô hình chuyên gia hỗn hợp; sử dụng lõi thưa thớt thế hệ thứ tư (SparseCore) để dỡ tải các thao tác giao tiếp tập thể nhằm che giấu độ trễ; tinh chỉnh phân bổ SRAM nhanh trên TPU (VMEM) để giảm tình trạng dừng bộ nhớ; và lựa chọn chiến lược phân mảnh tối ưu (như FSDP, TP, EP) dựa trên quy mô mô hình, kiến trúc và độ dài chuỗi. (Nguồn: InFoQ)