Представляємо Olmo-core 3: Відкрита, масштабована інфраструктура для навчання великих MoE-моделей
Сьогодні ми випускаємо Olmo-core 3 — значне оновлення нашого фреймворку для розробки великих мовних моделей із повністю переробленою відкритою системою навчання суміші експертів (MoE).
Olmo-core 3 розроблено для масштабування навчання MoE до діапазону трильйонів параметрів зі збереженням обчислювальної ефективності. Це одна з ключових систем наступного покоління Olmo, а також частина нашого постійного прагнення відкривати інструменти та інфраструктуру навчання, що лежать в основі кожної нової моделі.
Навчання великих моделей ШІ потребує значних обчислювальних ресурсів, підвищує витрати й споживання енергії та робить розробку передових моделей недоступною для багатьох академічних дослідників і невеликих лабораторій. Моделі MoE пропонують ефективніший підхід: вони можуть містити набагато більше навчених компонентів, або параметрів, не вимагаючи від кожного вхідного сигналу використовувати їх усі. Однак усю модель все одно потрібно зберігати в пам’яті GPU й оновлювати під час навчання, а спрямування вхідних даних до правильних експертів — спеціалізованих компонентів у MoE — у межах кластера створює власні витрати на комунікацію та координацію. У міру зростання MoE ці витрати можуть нівелювати значну частину обчислювальної переваги використання лише частини моделі для кожного вхідного сигналу.
Olmo-core 3 створено, щоб усунути цей розрив. В одному з тестів ми збільшили пул експертів із 8 до 128, водночас обираючи лише чотири експерти для кожного токена — невеликої одиниці тексту, яку обробляє мовна модель, — і зберегли приблизно сталу кількість активних параметрів на токен — близько 3,2 млрд. Загальна ємність параметрів зросла з 4,6 млрд до 47 млрд, тоді як пропускна здатність навчання знизилася менш ніж на 5%.
Цю саму інфраструктуру протестовано на загальній кількості параметрів понад один трильйон.
Побудова навчального стека з урахуванням реальної роботи MoE
Olmo-core розвивався з кожним поколінням Olmo.
Наші дослідження розріджених моделей розпочалися з OlmoE, яка використовувала архітектуру MoE із 64 маршрутизованими експертами. Натомість Olmo 3 використовувала щільну архітектуру, тобто майже вся модель була активною для кожного токена, а її навчальний стек будувався навколо цієї конструкції. Olmo-core 3 розширює фреймворк системою навчання, призначеною для значно більших моделей MoE.
Наша попередня реалізація MoE в Olmo-core використовувала повністю фрагментований паралелізм даних (FSDP), налаштований на збирання та повторне розподілення ваг моделі для кожного невеликого пакета навчальних даних. Olmo-core 3 переходить на систему, засновану на розподіленому паралелізмі даних (DDP). Вона зберігає експертів на GPU та спрямовує до них відповідні дані, уникаючи повторного збирання ваг.
Megatron-Core від NVIDIA є усталеним варіантом для навчання великих MoE. Olmo-core 3 додає інтегрований стек навчання MoE до фреймворку, що лежить в основі Olmo, із переробленою архітектурою, яка підвищує пропускну здатність порівняно з нашою попередньою реалізацією на основі FSDP. У попередньому тесті на восьми GPU NVIDIA B300 MoE із 47 млрд параметрів обробляла 52 000 токенів за секунду на GPU з новим стеком, порівняно з 19 400 у попередній реалізації — приблизно у 2,7 раза вища пропускна здатність.
Масштабування та оптимізація навчання MoE
Olmo-core 3 поєднує кілька методів розподілу великих MoE між кластерами GPU з оптимізаціями, які роблять маршрутизацію та обчислення ефективнішими.
Три методи визначають, як модель і її навчальний стан розподіляються між апаратними ресурсами:
- Паралелізм експертів розподіляє експертів між GPU, тому кожен GPU зберігає лише частину повного пулу експертів.
- Конвеєрний паралелізм розподіляє шари моделі — послідовні етапи, що перетворюють вхідні дані, — між групами GPU, зменшуючи обсяг моделі, який кожен GPU має зберігати в пам’яті.
- Розподілений оптимізатор розподіляє стан оптимізатора — додаткові дані, що використовуються для обчислення та застосування оновлень під час навчання, — між GPU замість зберігання повної копії на кожному GPU.
Разом ці методи дають змогу масштабувати MoE без потреби кожному GPU зберігати в пам’яті всю модель і її навчальний стан.
Olmo-core 3 також зменшує витрати на маршрутизацію даних до правильних експертів і виконання їхніх обчислень. Порозрядний паралелізм експертів розміщує маршрутизовані дані безпосередньо у вхідних буферах експертів, мінімізуючи додаткову роботу, необхідну для їх перестановки. Маршрутизація в GPU зберігає метадані маршрутизації на GPU, щоб CPU міг ставити завдання в чергу, не очікуючи копіювання цих даних назад. А згрупований GEMM об’єднує багато невеликих обчислень експертів, щоб GPU міг виконувати їх ефективніше.
Нарешті, Olmo-core 3 підтримує MXFP8 — формат чисел із нижчою точністю, який представляє деякі значення меншою кількістю бітів. Це може зменшити обсяг обчислень і даних, що передаються між GPU, якщо отримана економія перевищує витрати на перетворення між числовими форматами.
Ми виміряли вплив MXFP8 на наскрізну пропускну здатність навчання в контрольованому тесті на чотирьох GPU NVIDIA B300, рівномірно розподіливши роботу між експертами. Коли MXFP8 було ввімкнено для частин системи, де він давав найбільший ефект, пропускна здатність навчання була приблизно на 21% вищою, ніж із BF16 — форматом із вищою точністю, який ми використовували як базовий, — тоді як піковий обсяг активної пам’яті зменшився зі 103 ГіБ до 95 ГіБ. Більша частина приросту була зумовлена прямими обчисленнями та переміщенням даних між експертами, а не лише механізмом уваги.
Ці методи й оптимізації мають працювати разом. Прискорення однієї частини навчання може створити витрати в іншій; швидші обчислення можуть вимагати більшого переміщення даних, тоді як передавання меншої кількості бітів може не допомогти, якщо перетворення даних займає надто багато часу. Olmo-core 3 побудовано з урахуванням цих компромісів у всьому процесі навчання, що дає нам — і дослідникам, які використовують відкритий стек, — контроль над взаємодією його складових.
Перегляньте наш інтерактивний огляд, щоб дізнатися, як паралелізм даних, експертів і конвеєра працюють разом для масштабування навчання MoE — від одного GPU до багатьох.
Масштабування до діапазону трильйонів параметрів
Ми протестували Olmo-core 3 у різних конфігураціях на GPU NVIDIA B300, зокрема модель із 1,2 трлн параметрів, у якій на кожен токен активними були 58,36 млрд параметрів, розподілених між 512 GPU. Найвища зафіксована пропускна здатність становила 858 TFLOP/s/GPU — показник корисних обчислень моделі за секунду на кожному GPU. У цих тестах використовувалася випадкова маршрутизація для вимірювання продуктивності системи, а не якості навченої моделі.
Ми також експериментували з DeepEP v2 — альтернативним способом обробки комунікації між експертами на різних GPU — і досягли конфігурації із загальною кількістю параметрів 2,38 трлн. Це був короткий тест ємності, а не повний запуск навчання, тому він демонструє масштаб, якого може досягти Olmo-core 3, а не стабільну продуктивність під час навчання.
За таких масштабів продуктивність системи є лише частиною картини. У нашому технічному звіті також описано експерименти, які вплинули на те, як ми навчаємо MoE та вимірюємо їхню продуктивність. Наприклад:
- Оцінка, покликана заохочувати збалансовану маршрутизацію, могла покращуватися, навіть коли фактичне робоче навантаження ставало менш збалансованим. Цю помилку ми називаємо джеррімендерингом токенів.
- Зниження швидкості навчання експертів — величини їхніх оновлень під час навчання — через те, що вони обробляли менше токенів, не покращило результати в протестованому нами сімействі моделей.
- Обчислення на GPU займали різний час, коли змінювалися оброблювані значення, навіть за однакових розмірів матриць. Тому для порівняння продуктивності потрібні однакові вхідні значення, а не лише однакові форми.
- Перекриття комунікації та обчислень в окремих потоках GPU не завжди пришвидшувало навчання. У деяких тестах це сповільнювало наскрізне виконання — нагадування про те, що більше перекриття не обов’язково означає вищу пропускну здатність.
У звіті ці висновки пояснюються разом із підходами, які ми протестували та вирішили не застосовувати.
Створено для наступного покоління Olmo, відкрито для всіх
Olmo-core 3 є основою того, що ми створюємо далі. Olmo наступного покоління використовуватиме архітектуру MoE, і ми прагнемо, щоб це була найпотужніша Olmo, яку ми коли-небудь створювали, навчена на найбільшому наборі даних і з найдовшим контекстним вікном.
Новий стек дає нам змогу вийти за межі попередніх напрацювань із MoE, водночас забезпечуючи більшу гнучкість для адаптації навчання в міру розвитку моделей і апаратного забезпечення. Він також повністю відкритий: дослідники та розробники можуть використовувати Olmo-core 3 для навчання власних MoE, адаптувати його до різного апаратного забезпечення та експериментувати з маршрутизацією, паралелізмом та іншими частинами системи.
Це частина нашого підходу до розробки відкритих моделей: ваги моделі корисніші, коли інфраструктура та рішення щодо навчання, що лежать в їхній основі, також відкриті.
Щоб глибше ознайомитися з проєктуванням системи, експериментами, абляціями та підходами, які ми тестували на цьому шляху, прочитайте наш технічний звіт і дослідіть Olmo-core 3 на GitHub.
Спільнота
· Зареєструйтеся або увійдіть, щоб залишити коментар
Перекладено автоматично з англійської. Оригінал статті — за посиланням нижче.