Google phát hành hướng dẫn đào tạo nhà phát triển TPU Ironwood thế hệ thứ bảy, giải thích chi tiết tối ưu hóa hiệu suất cấp hệ thống

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