DeepSpeed налаштування: чому навчання впирається в пам'ять
Навчання великих моделей впирається не у швидкість прискорювачів, а в обсяг пам'яті. Ваги, градієнти й стани оптимізатора перебувають на пристрої одночасно, причому оптимізатор займає в рази більше, ніж сама модель.
Нестача пам'яті проявляється не на старті, а в середині прогону — коли вже витрачено години на підготовку даних і розігрів.
Ціна такого збою вища, ніж здається. Обладнання простоює між перезапусками, прогрес відкочується до останньої контрольної точки, а частину експериментів доводиться урізати: зменшувати розмір пакета, скорочувати довжину контексту, брати слабшу за потрібну модель.
Падіння процесу навчання гальмує не одну ітерацію, а весь напрям — від дослідження до виведення моделі в продукт.
Налаштування DeepSpeed — це керування тим, де саме живуть стани. Стадії ZeRO визначають, що шардиться між прискорювачами, вивантаження переносить частину стану в пам'ять хоста, а профіль конфігурації пов'язує вибрану схему з конкретним залізом і топологією обміну.
Нижче розбираємо, як це закривається на практиці — від вибору стадії до стабільного прогону без падінь і ручних перезапусків.
Що змінюється після грамотної конфігурації навчання
Без налаштування розподіленого навчання вузьке місце майже завжди одне: стани оптимізатора й активації не вміщаються в пам'ять прискорювача, тому інженер ріже розмір порції й втрачає завантаження. Пристрої простоюють в очікуванні синхронізації, і прогін розтягується.
Грамотна конфігурація розводить стани по пристроях і знімає це загальмування — та сама модель навчається в рази швидше.
DeepSpeed дає для цього набір режимів: розділення станів оптимізатора й градієнтів по пристроях, вивантаження ваг і активацій у пам'ять хоста, накопичення градієнтів дрібними порціями.
Сенс не у вмиканні всього одразу, а в підборі комбінації під конкретну модель і конфігурацію вузлів: надлишкове розділення додає обмін між пристроями й з'їдає виграш, недостатнє — повертає нестачу пам'яті.
Після підбору режимів завантаження обладнання виходить на рівні 85–95% замість провалів на синхронізації, а швидкість навчання перестає залежати від довжини послідовності.
Стійкість забезпечують контрольні точки: зберігаємо стан оптимізатора й планувальника, а не лише ваги, і пишемо їх асинхронно, щоб запис не гальмував лічбу.
Час відновлення після відмови вузла падає з годин до хвилин — процес піднімається з останньої контрольної точки, а не з нуля.
У сумі це дає високу частку прогонів, що доходять до цільової метрики без ручного втручання, і передбачувані строки: план навчання будується за виміряною швидкістю кроку, а не за оптимістичною оцінкою.
DeepSpeed налаштування: режими ZeRO і вивантаження пам'яті
Вибір режиму визначається тим, які стани навчання залишаються на кожному прискорювачі. Стани оптимізатора займають у рази більше місця, ніж самі параметри, слідом ідуть градієнти, і лише потім ваги.
Режими ZeRO відрізняються тим, яку з цих груп шардять між пристроями, а вивантаження пам'яті додає другий рівень — у пам'ять процесора або на локальний накопичувач.
| Варіант | Коли підходить | Обмеження |
|---|---|---|
| ZeRO-1: шардяться стани оптимізатора | Модель уже вміщається в пам'ять прискорювачів, але хочеться підняти розмір пакета даних | Параметри й градієнти дублюються на кожному пристрої — виграш за пам'яттю помірний |
| ZeRO-2: плюс градієнти | Навчання на кількох прискорювачах, коли обмін градієнтами вже став вузьким місцем | Зростає обсяг обміну між пристроями, чутливість до топології зв'язків |
| ZeRO-3: плюс параметри | Модель цілком не вміщається в пам'ять прискорювачів, потрібні довгий контекст і великий пакет | Параметри підтягуються на кожному кроці — вища затримка й вимоги до каналу обміну |
| Вивантаження станів у пам'ять процесора | Пам'ять прискорювачів у дефіциті, а крок чутливий до затримок | Обмежена пропускною здатністю шини між процесором і прискорювачами |
| Вивантаження станів на диск | Модель не вміщається навіть із повним шардуванням параметрів | Помітна просадка швидкості кроку, потрібен швидкий локальний накопичувач, зростає знос |
Як вибираємо на практиці
Спершу фіксуємо бюджет пам'яті прискорювачів і потрібний розмір пакета даних, потім знімаємо замір споживання на одному кроці навчання. Якщо ваги вміщаються — зупиняємося на першому або другому режимі, якщо ні — ідемо до повного шардування і лише потім до вивантаження.
Зміна режиму — це один параметр конфігурації, тому відкат до попереднього варіанта займає хвилини, а не переписування коду навчання.
Порядок впровадження: від аудиту навантаження до передачі
Перш ніж щось змінювати, потрібно зрозуміти, де саме навчання впирається: в обсяг пам'яті чи в обмін даними між прискорювачами. Тому робота починається із заміру поточного режиму, а не з правки конфігурації навмання — інакше ефект від змін не буде з чим порівняти.
Сім етапів
- Аудит навчання і профіль навантаження. Знімаємо профіль одного кроку: скільки пам'яті займають ваги, стани оптимізатора й активації, яку частку часу прискорювачі простоюють в очікуванні обміну. На виході — вузьке місце, назване числом, а не здогад.
- Схема розміщення. Підбираємо ступінь секціонування станів моделі (ZeRO) і режим вивантаження: що залишається на прискорювачі, що йде в пам'ять CPU, що перераховується на ходу. Схему фіксуємо до правок.
- Впровадження навчання. Змінюємо параметри по одному: розмір мікро-батчу, крок накопичення градієнта, перекриття обміну з обчисленнями, частоту контрольних точок. Після кожної зміни — замір.
- Серія контрольних прогонів. Короткі прогони на реальному наборі даних із звіркою кривої збіжності з еталонною. Мета — переконатися, що виграш у швидкості не куплений деградацією якості.
- Стабілізація довгих прогонів. Вмикаємо відновлення з контрольних точок, обмежуємо зростання їхньої кількості, перевіряємо поведінку при відмові одного процесу — навчання продовжується, а не падає цілком.
- Документування. Схема розміщення, значення параметрів і обґрунтування кожної зміни. Без цього конфігурація знову стане «чорною скринькою» після першої ж правки.
- Передача в експлуатацію. Проганяємо навчання разом із вашою командою, розбираємо типові збої й точки правки, залишаємо регламент запуску та зміни параметрів.
Що залишається в команди на виході
Передаємо не лише робочий запуск, а й усе, що потрібно, щоб команда відтворювала й супроводжувала його без нас. Конфігурація навчання оформлена як набір профілів під конкретний склад заліза, а не як вдалий підбір параметрів, що тримається в голові одного інженера.
- Робочі профілі запуску: кількість процесів, розмір порції даних на пристрій, крок накопичення градієнтів, ступінь розбиття станів моделі та режим винесення їх у пам'ять процесора або на диск.
- Схема розміщення моделі й станів: що живе на прискорювачах, що винесено за їхні межі, скільки займає кожен стан і за якого розміру моделі схема перестає вкладатися в ресурси.
- Регламент запуску: послідовність перевірок перед стартом — доступність усіх пристроїв, вільний обсяг пам'яті, цілісність останньої збереженої точки, збіг профілю з фактичним залізом.
- Регламент перезапуску: з якої точки продовжується навчання, як перевіряється збережений стан і за який час процес повертається в роботу після збою.
- Інструкція для чергової зміни: які записи в журналі читати, які ознаки деградації вважати аварійними, що виправляється на місці і коли викликається відповідальний.
- Журнал змін конфігурації: кожна правка з причиною та заміром до і після — час кроку, споживання пам'яті, динаміка якості на контрольній вибірці.
- Схема оповіщень: набір метрик і порогів, за якими зміна дізнається про проблему раніше, ніж навчання стане на добу.
Папка з профілями, схемами й регламентами лежить поряд із кодом і версіонується разом із ним. Приймання проводимо разом із вашою командою: черговий інженер самостійно проганяє запуск і відновлення на тестовому контурі, а ми відповідаємо на запитання щодо кожного пункту.
Кейс: як зняли обмеження за пам'яттю під час навчання
До нас прийшла команда, якій не вистачало пам'яті для наступної версії моделі.
Навчання падало не на першому кроці, а через кілька годин: заміри показали, що стани оптимізатора й градієнти займали більше пам'яті, ніж самі ваги, а частина прискорювачів простоювала в очікуванні обміну.
Що заміряли
Зняли профіль за кроками: пікове споживання пам'яті, утилізацію кожного прискорювача, частку часу на пересилання між картами й кількість прогонів, що доживали до кінця.
Це дало базу для порівняння — без неї будь-які зміни в режимі розміщення виглядають як віра, а не як інженерне рішення.
Що змінили і що отримали
Перевели навчання на DeepSpeed: увімкнули шардування станів оптимізатора, градієнтів і параметрів між картами (сімейство ZeRO), винесли частину станів на процесор і замінили зростання мікро-батчу на накопичення градієнта.
Активації перераховуємо на зворотному проході, а не тримаємо в пам'яті цілком — цей обмін обчислень на пам'ять тут вигідніший.
Розмір моделі виріс майже втричі на тій самій кількості прискорювачів. Частка прогонів, що доходили до кінця, піднялася приблизно з третини до всіх, збої за пам'яттю зникли як клас, а час одного кроку виріс на 10–15% — це усвідомлена плата за передбачуваність.
Контрольні точки зберігаються й відновлюються, відкат на попередню конфігурацію розміщення займає хвилини.
Чим ZeRO відрізняється від вивантаження станів на диск?
Це різні механізми, а не альтернативи. ZeRO ділить стани моделі — оптимізатор, градієнти, параметри — між прискорювачами, і кожен пристрій тримає лише свою частку. Вивантаження нічого не ділить: воно переносить стани в пам'ять процесора або на диск, звільняючи відеопам'ять.
Стадії ZeRO нарощують ефект: перша ділить стан оптимізатора, друга додає градієнти, третя — самі параметри. Плата — обмін між прискорювачами на кожному кроці: чим вища стадія, тим більший трафік і тим сильніше навчання залежить від пропускної здатності міжз'єднання.
| Варіант | Що знімає з прискорювача | Ціна | Коли підходить |
|---|---|---|---|
| ZeRO-1 | стан оптимізатора | обмін між прискорювачами | модель вміщається, але оптимізатор важкий |
| ZeRO-2 | оптимізатор і градієнти | трафік помітно зростає | потрібен запас під великі пакети |
| ZeRO-3 | плюс самі параметри | максимальний трафік | інакше модель не вміщається на групу |
| Вивантаження на CPU | оптимізатор, градієнти, буфери | пропускна здатність шини | пам'яті мало, час кроку терпить |
| Вивантаження на диск | стани й параметри | найповільніший крок | модель сильно перевищує сумарну пам'ять |
Вивантаження розміщує ті самі тензори в пам'яті процесора або на диску й підкачує їх до моменту обчислень. Пікове споживання відеопам'яті падає, але крок впирається у швидкість передавання по шині — виграш за пам'яттю оплачується часом кроку.
На практиці механізми комбінують: у DeepSpeed вони вмикаються незалежно, і зв'язка ZeRO-3 із вивантаженням параметрів навчає модель, яка інакше не вміщається на групу прискорювачів.
Ми рахуємо бюджет пам'яті під стан і активації, потім підбираємо режим і перевіряємо його на замірах реального кроку.
Часті технічні сумніви перед стартом
До старту інженерні команди перевіряють три речі: як налаштування ляже на наявний контур навчання, що буде зі зібраними чекпоінтами і як відкотитися, якщо результат не влаштує. Нижче — відповіді на те, що найчастіше обговорюємо на перших дзвінках.
Чи сумісно це з нашим контуром навчання?
Конфігурацію збираємо навколо вашого запуску: тренер, формат датасету й схема чекпоінтів залишаються незмінними. Змінюється профіль запуску — глибина ZeRO, вивантаження станів оптимізатора на CPU і диск, накопичення градієнтів, — а не код моделі. Перед масштабуванням звіряємо один крок навчання з поточним запуском.
Чи зросте час навчання?
Вивантаження міняє пам'ять на швидкість обміну, тому профіль підбираємо за замірами: скільки кроків за секунду дає кожен варіант і який розмір пакета вміщається.
Зазвичай звільнена пам'ять дозволяє підняти розмір пакета й кількість пристроїв і перекрити накладні витрати, але без замірів це не твердження.
Як відкотити зміни, якщо результат не влаштує?
Зміни ізольовані: конфігурація й обгортки запуску лежать окремо, попередній запуск залишається робочим.
Повернення — перемикання на старий профіль і продовження з останнього чекпоінта; формат чекпоінтів не змінюємо, тому проміжні результати читаються і попереднім, і новим запуском.
Як переконатися, що якість не впала?
Порівнюємо на короткому прогоні: ті самі дані й порядок, еталонний запуск і новий профіль, криві втрат і метрики на відкладеній вибірці. Розходження в межах шуму — продовжуємо; систематичний зсув — шукаємо причину в точності накопичення градієнтів і налаштуваннях оптимізатора.
Обговоримо задачу і повернемося з оцінкою та планом
План і оцінка з'являються не «на око», а з вхідних даних. Щоб порахувати етапи оптимізації навчання моделі, нам потрібно побачити поточну конфігурацію запуску, профіль навантаження на кілька прискорювачів і вимоги до часу кроку — без цього оцінка залишиться виделкою, а не планом.
Для старту достатньо:
- доступ до репозиторію з поточними скриптами запуску й конфігами шардингу та вивантаження — достатньо прав на читання;
- опис моделі та її розміру, кількість прискорювачів і тип зв'язку між ними;
- профіль реального навантаження: довжина послідовностей, частка дообучення, вимоги до пропускної здатності;
- журнали невдалих запусків або звіти про нестачу пам'яті, якщо вони збереглися;
- обмеження за середовищами: власна інфраструктура або хмара, вимоги до доступів та ізоляції;
- контакт інженера, який приймає результат, — щоб узгоджувати зміни без посередників.
Оцінка повертається документом: перелік етапів, очікуваний ефект кожного (утилізація прискорювачів, обсяг вивантажуваних станів, час кроку), ризики й точки відкату.
Перший етап починається не з правки конфігурації, а з відтворюваного заміру: фіксуємо базову пропускну здатність, знімаємо профіль пам'яті, і лише потім змінюємо стратегію шардингу та порядок вивантаження.
Опишіть задачу у вільній формі — повернемося з оцінкою та планом робіт.
