Google выпустила седьмое поколение руководства по обучению разработчиков TPU Ironwood, подробно о системных оптимизациях производительности

ME Новости: сообщение от 2 апреля (UTC+8). В последнее время официальный Google опубликовал для разработчиков учебное руководство по обучению, ориентированное на TPU седьмого поколения Ironwood. Это руководство направлено на то, чтобы помочь разработчикам в полной мере использовать системную производительность Ironwood TPU для эффективного обучения и развертывания передовых AI-моделей. Ironwood TPU — это настраиваемая AI-инфраструктура, разработанная для удовлетворения потребностей в вычислительной мощности для моделей с триллионами параметров; она с помощью технологий межчиповых соединений (ICI), оптических коммутаторов (OCS), сетей центров обработки данных (DCN) и объединения высокопропускной памяти (HBM) формирует полноценную систему, поддерживающую до 9,216 чипов. В тексте подробно описаны несколько ключевых стратегий оптимизации для данного оборудования, включая: использование его матричных блоков умножения (MXU), нативно поддерживающих FP8-обучение для повышения пропускной способности; применение JAX-библиотеки ядер Tokamax, специально оптимизированной для TPU, которая с помощью «разбрызгивающегося внимания» и «Megablox блочно-группового матричного умножения» обрабатывает неравномерные тензоры в моделях с длинным контекстом и смешанными экспертами; использование четвертого поколения разреженного ядра (SparseCore) для разгрузки операций коллективных коммуникаций с целью скрытия задержек; тонкую настройку распределения TPU-быстрой on-chip SRAM (VMEM), чтобы уменьшить задержки из-за памяти; а также выбор наилучшей стратегии разбиения (например, FSDP, TP, EP) в зависимости от масштаба модели, архитектуры и длины последовательности. (Источник: InFoQ)