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

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