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

ME News Новини, 2 квітня (UTC+8), офіційний сайт Google нещодавно опублікував керівництво для розробників щодо тренування восьмого покоління TPU Ironwood. Це керівництво спрямоване на допомогу розробникам у повному використанні системних можливостей TPU Ironwood для ефективного тренування та розгортання передових моделей штучного інтелекту. TPU Ironwood — це спеціалізована інфраструктура штучного інтелекту, створена для задоволення потреб у обчислювальній потужності моделей з трильйонами параметрів, яка за допомогою технологій міжчипового з’єднання (ICI), оптичних комутаторів (OCS), мережі дата-центрів (DCN) та високошвидкісної пам’яті з високою пропускною здатністю (HBM) побудувала повну систему, що підтримує до 9216 чипів. У статті детально описані ключові оптимізаційні стратегії для цього обладнання, зокрема: використання вбудованого підтримки тренування FP8 у матричних блоках (MXU) для підвищення пропускної здатності; застосування спеціалізованої бібліотеки ядра JAX, оптимізованої для TPU, Tokamax, для обробки довгих контекстів та нерегулярних тензорів у змішаних експертних моделях за допомогою “плескаючої уваги” та “Megablox групових матричних множень”; використання четвертого покоління розріджених ядер (SparseCore) для вивантаження операцій колективної комунікації з метою приховування затримок; тонке налаштування розподілу швидкодіючої внутрішньочипової SRAM-пам’яті (VMEM) TPU для зменшення простоїв пам’яті; а також вибір оптимальної стратегії розбиття (наприклад, FSDP, TP, EP) залежно від розміру моделі, архітектури та довжини послідовності. (Джерело: InFoQ)