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

Tin tức 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ố,
với 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 giới thiệu chi tiết các chiến lược tối ưu quan trọng dành cho phần cứng này, bao gồm:
tận dụng khả năng hỗ trợ đào tạo FP8 nguyên bản của đơn vị nhân ma trận (MXU) để tăng throughput;
sử dụng thư viện lõi JAX tối ưu 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 offload các thao tác giao tiếp tập thể;
tinh chỉnh 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)