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

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