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 News, ngày 2 tháng 4 (UTC+8), Google chính thức công bố 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 cho huấn luyện 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 các mô hình dài ngữ cảnh 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 chuyển tải các thao tác giao tiếp 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 (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)