← Усі статті

Три рівні стека мовної моделі: ядро, граф і система

Як влаштований програмний стек великої мовної моделі: обчислення, пам’ять, обмін даними і паралелізм навчання та виводу.

Три рівні стека мовної моделі: ядро, граф і система
Зміст

Коротко

Переклад на «Хабрі» збирає програмний стек великої мовної моделі в три рівні: ядро, граф і система. Цілі скрізь одні — більша пропускна здатність і менша затримка, — але впираються вони в три різні обмеження: обчислення, пам’ять і обмін даними. Автор не обіцяє повний огляд найкращих практик. Це мапа, за якою видно, на якому поверсі взагалі має сенс прискорювати модель.

Що сталося

Будь-яку таку систему проєктують під конкретні цілі і жорсткі межі. Для мовних моделей головні цілі — пропустити більше запитів і скоротити очікування відповіді. Межі задають швидкість і типи апаратних операцій, обсяг і ієрархія пам’яті, а також ширина, затримка і ієрархія передачі даних — і всередині машини, і мережею. Автор відсилає до розділу Roofline у Scaling Book від Google: там якраз видно, чи впираєтеся ви в рахунок, чи в підвіз даних.

Три рівні абстракції розводять ці межі по різних інструментах. Ядро — скалярні й векторні команди і команди над плитками даних; типові середовища — CUDA, C/C++, PTX, Triton. Граф — тензорні операції однієї моделі на одній відеокарті: тут важливі збираність і те, наскільки легко змінити схему. Система — як ділити пам’ять, обчислення і зв’язок між багатьма пристроями під час навчання і під час виводу. На кожному поверсі оптимізують своє, і прийом з нижнього рівня не замінює рішення на верхньому.

На рівні ядра показовий приклад — увага. Наївна формула змушує тягати великий рядок матриці уваги з глобальної пам’яті. Онлайновий softmax переписує ту саму формулу як рекурентну, і тоді не потрібні зайві читання і записи в глобальну пам’ять. На рівні графа історично зручні середовища програвали за швидкістю спеціалізованим: для серйозного виводу брали ONNX Runtime, TensorRT, FasterTransformer. Зараз PyTorch відіграв частку виводу: дешевшими стали графи CUDA і torch.compile, а самі мовні моделі настільки важкі, що накладні витрати каркаса вже не головна стаття. Типові прийоми графа — злиття сусідніх ядер, об’єднання операцій, квантування і розрідженість: менше ганяти дані і менше займати пам’ять за тих самих обчислень.

На системному рівні майже будь-який паралелізм піднімає пропускну здатність, але не будь-який знижує затримку. Паралелізм за даними ріже пакети: під час виводу обміну майже немає, під час навчання він слабкий, затримка не падає. Конвеєр ріже пакети і шари: зв’язок слабкий і під час навчання, і під час виводу, пам’ять під модель економиться, затримка знову не падає. Тензорний паралелізм ріже повнозв’язні шари за рядками і стовпцями і увагу за головами: затримка знижується, зате обмін дорогий. Вибір залежить від того, чого не вистачає — пам’яті чи каналу зв’язку. Обсяг обчислень при цьому зазвичай сталий; винятки на кшталт повністю розділеного паралелізму даних автор обмовляє окремо. Для зв’язки навчання і виводу він вказує на VeRL і стик з Megatron-LM та vLLM, а огляд стратегій — на документацію NeMo і DeepSpeed.

Чому це важливо

Без цієї мапи легко прискорювати не той поверх. Квантування не полагодить дорогий обмін за тензорного паралелізму, а зайва відеокарта не прибере затримку, якщо схема паралелізму її принципово не знижує. Інженеру, який обирає між «донавчити», «стиснути» і «розкласти на кілька карт», корисно спочатку назвати обмеження: рахунок, пам’ять чи зв’язок.

Другий практичний зсув — вивід більше не живе в окремому світі від навчання. PyTorch на виводі і спільні каркаси на кшталт зв’язки Megatron-LM з vLLM означають, що одна й та сама команда може говорити на одному стеку, а не тримати два несумісні ланцюжки. Це не скасовує спеціалізовані середовища там, де кожен відсоток затримки коштує грошей. Це прибирає звичку вважати «зручний граф» і «швидкий вивід» взаємовиключними.

На практиці

Варто читати чужий звіт про прискорення з питанням, на якому з трьох рівнів зроблено виграш. Якщо автор злив ядра, це граф однієї карти. Якщо він розклав голови уваги по пристроях, це вже система, і затримка з пропускною здатністю розійдуться. Плутанина цих поверхів — часта причина, чому «прискорення зі статті» не переноситься на ваш контур.

Таблиця паралелізму в перекладі якраз для такого вибору, а не для того, щоб увімкнути всі стратегії одразу. Пам’ять моделі і ціна обміну вирішують більше, ніж модність назви методу.

  1. Назвіть вузьке місце до вибору прийому: обчислення, пам’ять чи передача даних.
  2. На одній карті спочатку дивіться злиття операцій, квантування і те, чи не ганяє ядро зайві дані, як у прикладі з увагою.
  3. Якщо додаєте пристрої, перевірте, чи падає затримка: паралелізм за даними і конвеєр її не знижують.
  4. Не тримайте навчання і вивід у несумісних стеках без причини — подивіться, чи стикуються ваші каркаси так, як описано для Megatron-LM і vLLM.

Підсумок

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

Коментарі

Завантаження коментарів…