Передача ваги для RL післянавчання менш ніж за 2 секунди

Нещодавно ми досягли оновлення параметрів між машинами за 1,3 секунди для Kimi-K2 (1 трлн параметрів), передавши ваги з 256 навчальних графічних процесорів (BF16) на 128 графічних процесорів виведення (FP8).

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

Використовуючи одноранговий зв'язок RDMA, ми можемо зробити передачу ваги надзвичайно швидкою без зміни механізму виведення, а також спростити написання та підтримку коду.

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

python

def rdma_write(src_ptr, dst_ptr, size, src_mr, dst_mr):
    # Write from local [src_ptr, src_ptr+size) to remote [dst_ptr, dst_ptr+size).
    # src_mr and dst_mr contains Memory Region metadata of both sides.
    ...

Вузол призначення навіть не отримує сповіщення про передачу. Це забезпечує низьку затримку, високу пропускну здатність і передачу без копіювання, яку ініціюють навчальні вузли без будь-якої керуючої логіки на вузлах виведення.

Збір метаданих — контролер збирає метадані параметрів з усіх навчальних графічних процесорів і графічних процесорів виведення.

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

Розподіл розкладу — контролер надсилає розклад усім навчальним графічним процесорам.

Виконання — після кожного кроку навчання контролер подає сигнал навчальним графічним процесорам розпочати передачу.

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

Параметри під час навчання розподіляються відповідно до розміщення FSDP. Використовуючи full_tensor(), усі графічні процесори в DeviceMesh можуть реконструювати повний параметр, тому всі вони можуть слугувати джерелом для передачі ваги.

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

Конвеєр передавання ваг перекриває копіювання з хоста на пристрій, складання тензорів, ф’южн і квантування, а також передавання RDMA між MeshGroups.

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

Копіювання пам'яті з хоста на пристрій (memcpy) — якщо FSDP вивантажує вагу в ЦП.

Підготовка параметрів — реконструкція повної ваги за допомогою full_tensor(), застосування злиття проєкцій, квантування за потреби.

Передача RDMA — запис у пам'ять віддаленого графічного процесора виведення без копіювання

Глобальний бар'єр — після завершення всіх викликів full_tensor() виконується синхронізація між групами сіток за допомогою GLOO через Ethernet.

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

full_tensor() та інші операції графічного процесора створюють додаткове навантаження на пам'ять GPU. Щоб уникнути помилки нестачі пам'яті (OOM), ми починаємо виконання завдання лише тоді, коли поточні активні завдання займають менше тимчасової пам'яті GPU, ніж налаштоване порогове значення.

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

Передавання ваг на основі Rank0 спочатку збирає ваги, а потім розподіляє їх на вузол інференсу.

Точка-до-точки передавання переносить ваги з обох навчальних вузлів до мережевих інтерфейсів вузла інференсу.

Поширеним шаблоном є перенаправлення всіх параметрів через графічні процесори з рангом 0: збір на навчальному ранзі 0, надсилання на ранг 0 виведення, а потім повторний розподіл. Це швидко стає вузьким місцем, обмеженим пропускною здатністю PCIe та мережевою платою (NIC) одного графічного процесора (наприклад, 400 Гбіт/с ≈ 50 ГБ/с).

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

Передавання на основі RPC використовує буфер на боці інференсу та виклики worker, тоді як одностороннє передавання записує ваги безпосередньо, без участі worker інференсу.

Деякі системи покладаються на виклик методу update_weight() механізму виведення для кожного тензора. Це означає вторгнення в код виведення, а також накладні витрати від RPC, серіалізації та координації площини керування.

Завдяки примітиву RDMA WRITE ми оновлюємо ваги непомітно в пам'яті графічного процесора виведення без додаткових копій. Не задіюються жодні повідомлення площини керування чи логіка керування ЦП. Жодних модифікацій механізму виведення не потрібно.

Процес передачі ваги може використовувати чотири типи апаратних ресурсів: (1) переміщення даних між хостом і пристроєм, (2) обчислення на графічному процесорі для злиття проєкцій та квантування, (3) мережу RDMA для площини даних, (4) Ethernet для площини керування.

Наша конструкція розбиває завдання передачі ваги на етапи конвеєра, що дозволяє легко перекривати різні апаратні ресурси.

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

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

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

У нашій інженерній практиці ми розділяємо ці кроки на окремі компоненти. Кожен компонент можна протестувати окремо, проаналізувати та оптимізувати ізольовано.

Швидка та надійна передача ваги є важливим будівельним блоком для великомасштабного донавчання за допомогою підкріплення. Поєднавши примітив RDMA WRITE, статичний розклад передачі та конвеєрне виконання, ми скоротили оновлення трильйонів параметрів у Kimi-K2 лише до 1,3 секунди. Цей підхід простий для розуміння, легкий в обслуговуванні та дозволяє уникнути вузьких місць традиційних архітектур.

BibTeX

@online{perplexity2025WeightTransferForRlPostTrainingInUnder2Seconds,
  author = {{Perplexity} and Lequn Chen and Erin McKnight and Vivian Wang},
  title = {{Weight transfer for RL post-training in under 2 seconds}},
  date = {2025-09-24},
  url = {https://www.perplexity.ai/hub/blog/weight-transfer-for-rl-post-training-in-under-2-seconds}
}

RDMA WRITE: однобічна передача

Загальний робочий процес

Виконання передачі ваги

DeviceMesh і групи сіток

Конвеєр завдань

Керування використанням пам'яті графічного процесора

Чому це швидко та просто

Точковий зв'язок

Однобічна передача даних

Конвеєризація

Статичний розклад

Чітке розділення

Висновок

Цитування