Налаштування FSDP (Fully Sharded Data Parallel) для навчання

FSDP PyTorch: чому навчання великої моделі впирається в пам'ять

Навчання великої моделі впирається не в саме число параметрів, а в те, що кожен прискорювач зобов'язаний тримати в себе. Повний стан моделі — ваги, градієнти та стан оптимізатора — під час розподіленого навчання дублюється на кожному пристрої.

Стан оптимізатора при цьому в кілька разів важчий за самі ваги: на кожен параметр зберігається кілька накопичених величин.

Щойно сумарний обсяг перестає вміщатися в пам'ять прискорювача, навчання або падає з помилкою виділення, або починається вимушена економія: зменшують розмір пакета, вмикають накопичення градієнтів, ділять шари вручну.

Кожен такий крок — компроміс: вільної пам'яті більше, але час кроку зростає, а частина пристроїв простоює в очікуванні.

FSDP PyTorch закриває це шардуванням стану моделі: ваги, градієнти та стан оптимізатора розбиваються між пристроями, а не зберігаються повністю на кожному.

Перед обчисленням шару його параметри збираються на потрібних прискорювачах, після — знову розпускаються, а стан оптимізатора залишається розбитим постійно.

Саме там звільняється основна частина пам'яті, і навчання великих моделей перестає залежати від того, чи вміщається модель повністю в один прискорювач.

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

Що дає шардування стану моделі та оптимізатора

Шардування стану моделі та оптимізатора розподіляє параметри, градієнти й моменти по пристроях замість дублювання на кожному. Навчання зазвичай впирається не в саму модель, а в її оптимізатор: ваги ще вміщаються, а стани — вже ні.

У PyTorch FSDP це вмикається режимом шардування, і далі результат залежить від того, як налаштовані обмін і контрольні точки.

Перевіряємо за пунктами, що змінюється вимірювано.

• Розмір пакета. Звільнена пам'ять іде під активації, тому глобальний розмір пакета зростає без зміни обладнання. Більший пакет дає рівніший градієнт і дозволяє підвищити швидкість навчання без втрати збіжності.
• Час кроку навчання. Обмін градієнтами перекривається обчисленнями, тому крок подовжується помітно повільніше, ніж зростає обсяг даних. Підсумковий час навчання скорочується кратно, а не на відсотки.
• Стійкість до збоїв. Шардований стан збирається в контрольну точку без дублювання, тому падіння окремого процесу не означає втрату прогресу. Зникає клас відмов, пов'язаних із переповненням пам'яті на одному пристрої.
• Час відновлення. Запуск з останньої контрольної точки займає хвилини: стан піднімається в тому самому розбитті, і навчання продовжується з того самого кроку.

Перевіряти це варто на своєму профілі навантаження — у різних архітектур вузьке місце зміщується між пам'яттю оптимізатора, активаціями та мережею обміну.

Ми знімаємо заміри до і після, фіксуємо розмір пакета, час кроку та час відновлення, і тільки потім вважаємо режим налаштованим.

Порівняння режимів шардування та їхніх обмежень

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

Групи параметрів дозволяють точково виводити з шардування окремі шари — ембединги, нормалізації, вихідну голову — коли їхнє дроблення дає більше обміну, ніж економії. Нижче — три режими та межі, за якими кожен програє.

Режим закріплюють за замірами, а не за описом: дивляться частку часу в обміні та утилізацію пам'яті на робочому розмірі пакета. У PyTorch перемикання між режимами — це зміна конфігурації обгортки, саму модель переписувати не потрібно.

Як проходить впровадження: від аудиту до передачі команді

Навчання великої моделі впирається в пам'ять одного прискорювача, тому стан моделі та оптимізатора розподіляють по кількох.

На цьому переході спливають помилки конфігурації: невідповідність режимів точності, зайві пересилання між пристроями, розходження збіжності з попереднім запуском. Впровадження ведемо короткими кроками, де кожен підтверджується вимірюванням.

1. Аудит конфігурації навчання. Розбираємо поточний рецепт: розмір батчу та його розбиття, режим точності, склад шарів, планувальник швидкості навчання. Фіксуємо базову лінію — час кроку, споживання пам'яті, криву збіжності на короткому прогоні.

2. План шардування. Вирішуємо, що розносити по пристроях: параметри, градієнти, стан оптимізатора, а що залишити реплікованим. Визначаємо, які блоки обгортати повністю, щоб не платити за зайвий обмін.

3. Налаштування FSDP. Збираємо конфігурацію під вибраний план: політика обгортки, порядок перебудови параметрів, накопичення градієнтів. Перевіряємо, що запуск відтворюється на тому самому числі пристроїв.

4. Перевірка збіжності на реальній задачі. Порівнюємо з базовою лінією за кроками, а не лише за фінальною метрикою: розходження на початку вказує на точність або порядок обміну. Правимо до виходу на повний прогін.

5. Контрольні точки. Налаштовуємо збереження й завантаження розподіленого стану, відновлення після збою вузла та ротацію: прогін продовжується з останньої точки, а не з нуля.

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

Конфігурації, схеми та регламенти, які отримує замовник

Усе, що залишається після впровадження, лежить у вашому репозиторії та в документації, а не в листуванні й не в голові окремого інженера. Людина, яка прийде на проєкт після нас, розгортає навчання за тими самими файлами й отримує той самий результат. Нижче — що саме фіксується й передається.

• Конфігурації запуску — параметри навчання та розбивки моделі в системі контролю версій: зміни проходять погодження, а не правляться прямо на сервері, тому будь-який запуск відтворюваний.
• Схеми шардування — як FSDP у PyTorch розкладає параметри й стани оптимізатора по пристроях і скільки пам'яті резервує, плюс розрахунок під вашу модель і число вузлів.
• Політики доступу — хто має право запускати навчання, змінювати конфігурації та вивантажувати ваги; ролі розведені, доступ до даних і артефактів мінімально необхідний.
• Регламент відновлення — періодичність контрольних точок, строк зберігання та порядок перезапуску з останньої збереженої після відмови вузла.
• Документація — розгортання з нуля, опис конфігурацій і порогів, за яких навчання деградує, розбір типових помилок.
• Спостережуваність — криві втрат і пропускної здатності за кроками: одразу видно, чи впирається навчання в дані, мережу або прискорювач.

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

Кейс: навчання великої моделі перестало впиратися в пам'ять

Задача була типовою для великих моделей: навчання не вміщалося в пам'ять прискорювачів за осмисленого розміру пакета. Команда йшла на компроміси — зменшувала пакет, урізала довжину контексту, ділила шари вручну, — але кожен другий запуск усе одно падав з помилкою нестачі пам'яті. Інженерні години йшли не на якість моделі, а на перезапуски й підбір конфігурації.

Що було до

Базовий замір на реальній задачі: крок навчання займав близько 3,8 секунди, успішно доходили до контрольної точки 6 запусків із 10, решта падали через пам'ять.

За робочий тиждень набиралося 12–14 таких відмов, частина прискорювачів простоювала, а відновлення після збою потребувало ручного складання стану. Збільшення пакета, яке покращило б збіжність, було неможливим у принципі.

Що змінили

Навчання перевели на FSDP у PyTorch — параметри, градієнти та стани оптимізатора шардуються по пристроях, тому повна копія моделі більше не зобов'язана жити в пам'яті одного прискорювача.

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

Що стало

Крок навчання скоротився до 2,1 секунди, частка успішних запусків зросла до 10 із 10, падінь через пам'ять немає вже місяць. Розмір пакета вдалося підняти вдвічі — збіжність стала рівнішою, а відкат на попередню контрольну точку займає хвилини замість ручного перескладання. Ми замірюємо такі метрики на кожному проєкті: вони показують, чи окупилося налаштування розподілу.

Чим відрізняється FSDP PyTorch від звичайного розподіленого навчання?

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

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

Паралелізм моделі розв'язує задачу інакше: шари або блоки розрізаються між пристроями, і кожне тримає лише свою частину ваг. Плата — ручне розбиття, балансування навантаження й постійні пересилання активацій між ступенями.

Такий підхід масштабує обсяг моделі, але погано переноситься між архітектурами й потребує окремого налагодження під кожну.

FSDP PyTorch займає проміжне положення. Параметри, градієнти та стани оптимізатора шардуються по всіх процесах групи: перед обчисленням шар збирається повністю, після — звільняється.

На кожному прискорювачі постійно живе лише частка стану моделі, а не повна копія, при цьому цикл навчання залишається звичним — модулі обгортаються, далі робота йде як зі звичайною моделлю.

Звідси й критерій, коли застосовувати: мала й середня модель — паралелізм даних, менше обмінів і простіше налагодження; модель не вміщається в пам'ять одного прискорювача навіть із накопиченням градієнтів — шардування стану.

Для надглибоких мереж має сенс комбінація: FSDP усередині вузла, паралелізм даних між вузлами. Різниця підходів зводиться до одного питання — масштабуємо пропускну здатність реплік, обсяг стану чи і те й інше одразу.

Технічні сумніви перед запуском шардованого навчання

Сумніви інженерного рівня зазвичай ті самі: чи стане FSDP на наявний код моделі, чи витримають зв'язки між вузлами обмін градієнтами, що робити з накопиченими контрольними точками й чи не просяде якість за змішаної точності. Розбираємо по порядку.

FSDP у PyTorch не вимагає переписувати архітектуру: він підключається політикою обгортки поверх модулів, а оптимізатор і цикл навчання залишаються на місці.

Окремо перевіряємо вимоги до обладнання: межею масштабування частіше стає міжвузловий обмін, а не обсяг пам'яті на пристрої.

Що робити з накопиченими контрольними точками?

Старі файли залишаються валідними: розбираємо їх за шарами й збираємо в шардований формат окремим скриптом, зворотна конвертація дає відкат на нерозділений запуск. Перевірка — звірка виходу моделі до і після конвертації на одному кроці.

Чи вистачить зв'язків між вузлами й як бути зі змішаною точністю?

Замірюємо частку часу на обмін між вузлами: якщо вона зростає, збільшуємо розмір шарда й скорочуємо число синхронізацій. Змішана точність вмикається на рівні обгорток; ризик у перерахунку нормалізації знімається приведенням типів на межах блоків і коротким прогоном.

Обговоримо задачу: що уточнити перед стартом робіт

Чим конкретніші вхідні дані щодо задачі, тим менше припущень в оцінці. Нам не потрібен доступ до продакшену на цьому етапі — достатньо відповідей на запитання нижче, щоб зрозуміти обсяг робіт і ризики.

• Поточний цикл навчання: на чому запускається зараз, які вузькі місця вже видно — нестача пам'яті на пристрої, простій прискорювачів, довгі паузи на збереженні стану.
• Розмір моделі та склад пам'яті: скільки параметрів, які шари найважчі, що саме не вміщається — ваги, градієнти, стани оптимізатора або проміжні активації.
• Конфігурація обладнання: число прискорювачів, як вони пов'язані між собою, пропускна здатність каналу зв'язку — від цього залежить вибір схеми шардування і те, чи окупиться вона.
• Схема розміщення: що розкладаємо по пристроях, які блоки залишаємо повністю на кожному, чи потрібна перебудова при зміні числа вузлів.
• Робота з контрольними точками: у якому вигляді зберігаємо стан, чи потрібне відновлення навчання з тієї самої точки після збою і як швидко воно має відбуватися.
• Сценарій відмови: що вважаємо допустимим — втрата кроку, перезапуск з останньої контрольної точки, відновлення на меншому числі пристроїв.
• Критерії приймання: за якими метриками оцінюємо результат — час кроку, прискорення при зростанні числа пристроїв, збіжність і підсумкова якість моделі.

Далі ми повертаємося з оцінкою та планом: обсяг робіт, строки, склад команди й точки контролю.

Для FSDP у PyTorch окремим пунктом пропонуємо пілотний запуск на вашій моделі та реальній ділянці даних — він показує фактичне прискорення й знімає питання щодо сумісності до початку основних робіт.

Опишіть задачу в тому вигляді, у якому вона є — чернеткою, без підготовки. Уточнювальні запитання поставимо самі, а строки робіт і межі відповідальності зафіксуємо до старту.