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

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