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

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