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

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