HeyGen прискорила Avatar IV на TPU у 1,86 раза: як оптимізували генерацію потокового відео

HeyGen та Google Cloud детально описали перенесення Avatar IV на восьмичиповий TPU Trillium v6e. Модель створює talking-head відео з фотографії та аудіо, тому кожен фрагмент має бути готовий до жорсткого дедлайну: якщо один чанк запізнюється, потокове відтворення зупиняється. Команда не просто змінила тип прискорювача, а перебудувала критичні ділянки обчислювального конвеєра.

Як влаштований конвеєр Avatar IV

Avatar IV перетворює одну фотографію та аудіодоріжку на рухоме відео людини. Усередині послідовно працюють diffusion transformer, який генерує рух відповідно до звуку, другий transformer для підвищення роздільності та VAE-декодер, що перетворює латентне представлення на пікселі. Відео формується у 720p або 1080p зі швидкістю 25 кадрів за секунду й віддається частинами ще до завершення всієї генерації.

Два трансформери разом мають понад 36 ГБ ваг у bf16, тоді як один чип Trillium містить 32 ГБ HBM. Тому ваги розділили між вісьмома чипами за FSDP, а токени відеопослідовності — за Ulysses. Виробничий PyTorch-код залишився майже незмінним завдяки torchax: операції потрапляють у JAX і компілюються XLA, а TPU-специфічна оптимізація зосереджується у власних Pallas kernels.

Шість етапів прискорення

Команда вимірювала не швидкість окремого kernel benchmark, а повний час створення чанка. Після шести наборів змін тривалість скоротилася трохи більше ніж удвічі від першої робочої TPU-версії, що дало підсумкове прискорення 1,86×. Найбільший перший крок забезпечили спеціалізовані attention kernels, правильне розбиття послідовності, налаштування XLA і tile sizes під конкретні форми тензорів.

За даними HeyGen, фінальний конвеєр показав швидкість, порівнянну з виробничою конфігурацією 8×H100, і міг бути до 25% дешевшим у перерахунку на хвилину створеного відео. Ці цифри належать конкретній моделі та формам тензорів, але сам метод вимірювання корисний для будь-якого real-time inference: оптимізація зараховується лише тоді, коли виграє весь шлях до готового результату.

Як сховали all-to-all обміни

Ulysses sequence parallelism виконує пару all-to-all операцій усередині кожного self-attention. Спочатку вони стояли на критичному шляху: передавання даних уже використовувало 85–90% пропускної здатності mesh, а монолітна послідовність «обмін — attention — обмін» не залишала XLA незалежної роботи, за якою можна приховати комунікацію.

Голови attention розділили на незалежні групи. Поки одна група передає дані, інша виконує матричні обчислення. Сам обсяг трафіку не зменшився, але більша його частина перестала впливати на час чанка. У production traces видимий слід колективної операції на compute stream скоротився приблизно у п’ять разів.

Маску sparse attention прибрали конструктивно

Найбільший kernel етапу super-resolution працює з десятками тисяч токенів. Кожен кадр бачить вікно сусідніх кадрів і один глобальний reference frame. Стандартні блоки по 128 токенів не збігалися з межами кадрів, тому приблизно кожен п’ятий блок був частковим і вимагав перевірок mask predicates, padding та додаткового проходу.

Замість прискорення складної mask-логіки розмір блока дозволили зменшувати з кроком 16 — найдрібнішим bf16 tile, який підтримує апаратне забезпечення. Блоки почали точно ділити простір одного кадру й стали або повністю активними, або повністю пропущеними. Це прибрало часткові блоки, padding та другий attention pass.

Як розірвали послідовний ланцюг softmax

Flash-style attention зазвичай підтримує поточний максимум для кожного рядка запиту й перемасштабує накопичувач, коли наступний блок ключів підвищує його. Через це кожен блок залежить від попереднього. Команда замінила running maximum на верхню межу, яку можна заздалегідь обчислити через нерівність Коші—Шварца за нормою query і найбільшою нормою key.

Попередньо розрахований невеликий масив норм передається kernel через scalar prefetch. Для 98–99% attention heads межа виявилася достатньо точною, а решта автоматично повертається до стандартного online softmax. Після усунення залежності довелося повторно підібрати геометрію блоків, адже оптимальна конфігурація змінилася.

Layout став частиною контракту з компілятором

Нормалізація, rotary embeddings, projections і пакування голів створювали дані у layout, який міг не відповідати фізичному формату наступної all-to-all операції. XLA додавав ланцюжок копіювань і repack. Команда об’єднала підготовчі операції в один Pallas kernel, який відразу записує результат у формат вхідного буфера collective.

Подібним контрактом стали scheduler flags і чесна оцінка вартості власних kernels. Якщо компілятор вважає custom kernel операцією з нульовими FLOPs та байтами, latency-hiding scheduler неправильно оцінює, що можна перекрити. Після додавання реалістичних cost estimates розклад покращився без зміни самої моделі.

Якість пікселів перевіряли на двох рівнях

Оптимізації першого рівня мали давати byte-identical відео: хеш кожного кадру повинен збігатися з baseline. Зміни, що переставляють порядок редукцій і неминуче створюють невеликі bf16-відмінності, проходили другий quality gate — вузький діапазон числової схожості та сліпу покадрову перевірку власниками моделі.

Один із варіантів із перенесенням residual stream у bf16 справді був швидшим, але не пройшов межу якості, тому його відхилили. Це важлива деталь: команда не оголосила будь-яке зменшення latency успіхом, якщо воно було отримане ціною помітної зміни відео.

Висновок редакції

Кейс Avatar IV показує, що продуктивність великої відеомоделі визначає не лише швидкість прискорювача. Вирішальними стали розклад колективних операцій, відповідність блоків структурі даних, математичне усунення послідовної залежності та явні контракти з компілятором. Для українських AI-команд практичний урок полягає у повному трасуванні конвеєра, окремих quality gates і перевірці кожного прискорення на кінцевому результаті, а не на ізольованому тесті.