Як Mira прискорює MoE-моделі на одному GPU

Чому MoE-моделі перевантажують пам'ять

Mixture-of-Experts розділяє параметри на безліч малих підмереж — експертів. Для кожного токена активується лише кілька з них. Така архітектура дозволяє масштабувати загальну кількість параметрів без зростання обчислень на токен. Виглядає привабливо для розгортання на одному GPU: активних параметрів мало. Але всі експерти мають десь зберігатися. Неактивні вони все одно займають більшість пам'яті. Маршрутизація на рівні токенів вирішує, які експерти спрацюють. Рішення динамічне й непередбачуване. Крім того, воно перекошене: деякі експерти викликають набагато частіше інших. Існуючі системи реагують на маршрутизатор: чекають виводу, потім переносять обраних експертів із повільної пам'яті у VRAM. Реактивний підхід марнує кеш. Він також забороняє накладати передачі даних на обчислення, коли VRAM обмежений. GPU простаює, поки експерти подорожують шиною.

Як Mira передбачає і попередньо завантажує експертів

Mira змінює підхід з реактивного на проактивний. Додає легкі предиктори для кожного шару. Кожен предиктор аналізує поточний маршрут і прогнозує, яких експертів знадобляться наступні два шари. Прогноз прибуває раніше, ніж спрацюють відповідні маршрутизатори. Цей часовий запас дозволяє почати переміщення параметрів заздалегідь. Предиктори настільки малі, що накладні витрати незначні.

Вони наповнюють дворівневий кеш на GPU. Перший рівень — HOT — тримає експертів, яких телеметрія маршрутизації показує як часто використовувані. Другий — STAGE — зберігає тих, кого предиктори вважають майбутніми. Токенна телеметрія керує обома рівнями. Вона відстежує реальне використання, щоб HOT залишався точним. Вона також перевіряє прогнози, щоб не захаращувати STAGE помилками. Коли передбачення не вдається, система швидко витісняє помилково завантаженого експерта. Такий дизайн тримає кеш зосередженим на реальних потребах моделі.

Дворівневий кеш і власна стиснення

Переміщення параметрів експертів вимагає пропускної здатності. Mira скорочує витрати власною формат стиснення, створеним для тензорів експертів. Формат зменшує метадані і пакує параметри ефективніше за універсальні кодеки. Автори стверджують: втрата точності мінімальна. Стиснені експерти швидше подорожують між хост-пам'яттю і GPU. Середовище виконання координує предиктори та політики витіснення кешу. Окрім того, воно керує квантованими передачами. Воно планує передачі так, щоб вони перекривалися з обчислювальними ядрами. Перекриття — ключ до приховування затримок при обмеженому VRAM. Стаття описує середовище як повністю інтегроване. Зовнішній оркестратор не потрібен. Система сама вирішує, коли попередньо завантажувати, що витісняти і як упаковувати кожну передачу — без налаштувань користувача.

Виміряні прискорення на обмеженому залізі

Автори порівняли Mira з передовими базовими лініями