Самая быстрая в мире синхронизация дельт весов

В нашем предыдущем техническом блоге мы показали синхронизацию дельт весов для чекпоинта GLM-5.3 объемом ~770 ГБ за 36 секунд с паузой запросов в 6 секунд. Теперь эта же синхронизация занимает 15 секунд при паузе в 0,8 секунды, обеспечивая рост пропускной способности до 3,2 раза.

Мы перестроили процесс синхронизации весов вокруг тех весов, которые уже находятся на GPU. Это избавило нас от необходимости поддерживать вторую полную копию в памяти CPU. Такой новый дизайн позволяет нам:

• Проверять обновления, считывая только измененные области памяти GPU

Проверять обновления, считывая только измененные области памяти GPU

• Минимизировать время паузы за счет предварительной подготовки записей на GPU

Минимизировать время паузы за счет предварительной подготовки записей на GPU

• Использовать освобожденную хост-память для кеша KV, повышая пропускную способность

Использовать освобожденную хост-память для кеша KV, повышая пропускную способность

• Запустить обучение с подкреплением (RL rollout) для Kimi K3 на 8:GB300, чему раньше мешали ограничения памяти CPU

Запустить обучение с подкреплением (RL rollout) для Kimi K3 на 8:GB300, чему раньше мешали ограничения памяти CPU

Краткий обзор синхронизации дельт весов

Трейнер в окружении клиента публикует полные чекпоинты и разреженные дельты в объектное хранилище. Каждый тензор дельты представляет собой zstd(parent_bytes XOR new_bytes). Как только снимок загружен, трейнер отправляет запрос на синхронизацию дельты весов с указанием обслуживаемой политики. Кластеры инференса по всему нашему парку загружают каждый объект один раз через Baseten Delivery Network (BDN), восстанавливают новую политику и обновляют свои GPU на месте. Важно также отметить, что трейнер и GPU инференса не разделяют никакой инфраструктуры.

Эта часть осталась без изменений.

Рисунок 1: Трейнер публикует каждый чекпоинт и дельту один раз в объектное хранилище в своем собственном окружении; каждая реплика вытягивает одни и те же объекты через Baseten Delivery Network и производит синхронизацию самостоятельно.

Копия в памяти CPU была узким местом

Наша первая архитектура хранила полный родительский чекпоинт в оперативной памяти хоста. Декодеру он был необходим для преобразования байтов XOR в новые значения, а после каждого обновления его приходилось исправлять, чтобы следующая дельта декодировалась относительно правильного родителя.

Эта копия обходилась нам дорого в трех аспектах:

• Память. Каждая реплика хранила своего полного родителя в оперативной памяти хоста: около 770 ГБ для GLM-5.3. На узле 8×B300 эта копия и кеш KV на хосте объемом 1,1 ТБ не могли сосуществовать; кеш пришлось уменьшить до 275 ГБ.

Память. Каждая реплика хранила своего полного родителя в оперативной памяти хоста: около 770 ГБ для GLM-5.3. На узле 8×B300 эта копия и кеш KV на хосте объемом 1,1 ТБ не могли сосуществовать; кеш пришлось уменьшить до 275 ГБ.

• Время. Перевод копии в CPU на новую политику занимал около 45 секунд на одно обновление, блокируя следующую дельту.

Время. Перевод копии в CPU на новую политику занимал около 45 секунд на одно обновление, блокируя следующую дельту.

• Два источника правды. Каждое обновление должно было попадать и на GPU, и в копию на CPU, иначе следующая дельта декодировалась бы относительно неправильного родителя.

Два источника правды. Каждое обновление должно было попадать и на GPU, и в копию на CPU, иначе следующая дельта декодировалась бы относительно неправильного родителя.

GPU уже содержит каждый байт политики N. Нам нужен был лишь способ читать те немногие байты, которые изменяются.

Рисунок 2: Раньше каждая дельта декодировалась относительно копии на CPU размером ~770 ГБ, перевод которой занимал 45 секунд; теперь родительскими являются весы на GPU, а то же обновление применяется за 15 секунд при паузе в 0,8 секунды.

GPU является родителем

При запуске мы загружаем модель с локального диска, проверяем каждый обслуживаемый тензор по нему через актуальную структуру GPU движка, записываем контрольную сумму Adler-32 для каждого тензора и удаляем копию. С этого момента память GPU хранит единственного родителя.

Когда приходит дельта, наш нативный декодер напрямую обрабатывает команды zstd и сохраняет только ненулевые байты XOR: для каждого измененного элемента — его позицию и маску XOR. Он никогда не материализует плотный тензор XOR и никогда не читает веса модели.

Затем, пока старая политика продолжает обслуживание, каждый воркер GPU:

• Собирает старые значения только в измененных позициях.

Собирает старые значения только в измененных позициях.

• Применяет маску XOR для получения абсолютных новых значений.

Применяет маску XOR для получения абсолютных новых значений.

• Проверяет результат по контрольной сумме трейнера.

Проверяет результат по контрольной сумме трейнера.

• Записывает новые значения в план записи в нативном формате тензоров движка.

Записывает новые значения в план записи в нативном формате тензоров движка.

Рисунок 3: Все действия вплоть до записи выполняются во время работы старой политики; во время паузы каждый ранг записывает абсолютные значения, и все ранги должны прийти к согласию до того, как новая политика будет опубликована.

Контрольные суммы дельт с обновлениями O(1)

Трейнер публикует контрольную сумму Adler-32 для новых байтов каждого измененного тензора. Наивная проверка этого означает повторное хеширование всего тензора, что неэффективно.

Каждый байт вносит вклад в Adler-32 независимо, величина которого зависит только от его (позиции, значения). Таким образом, влияние измененного байта на контрольную сумму может быть вычислено за константное время на основе его позиции, старого и нового значений. Мы начинаем с контрольной суммы, записанной при запуске, добавляем по одной поправке на каждый измененный байт и сравниваем с данными трейнера. Неизмененные байты никогда не читаются.

Эта оптимизация сделала все обновление до 2,3 раза быстрее (Kimi K3 при дельте 30 ГБ ускорился с ~230 до ~100 секунд), обеспечивая при этом корректность. Поврежденная загрузка или неверный родитель отклоняются до любой записи на GPU.

Рисунок 4: Читается только измененный байт: его старое и новое значения корректируют сохраненную контрольную сумму за константное время, и исправленная контрольная сумма должна совпадать с суммой трейнера.

Пауза длительностью менее секунды

Цель столь масштабной подготовки перед паузой — минимизировать операции, выполняемые во время самой паузы. Теперь пауза выглядит так:

• Закрыть прием новых запросов и поставить выполняющиеся запросы на паузу на месте.

Закрыть прием новых запросов и поставить выполняющиеся запросы на паузу на месте.

• Записать новые значения на месте + обновить производные веса.

Записать новые значения на месте + обновить производные веса.

• Опубликовать новую версию политики, возобновить прием запросов.

Опубликовать новую версию политики, возобновить прием запросов.

Для GLM-5.3 пауза теперь составляет 0,8 секунды против прежних 6. Кроме того, исчезло 45-секундное обновление на CPU, поэтому следующая дельта может примениться сразу же после поступления.

Дальнейшие оптимизации

Отказ от чекпоинта на CPU стал огромным прорывом. Тем не менее, без копии в CPU работа по обновлению делит ресурсы GPU и CPU, обслуживающих трафик. Эти оптимизации не дают ей снизить пропускную способность инференса.

Декодер дельт с учетом параллелизма

При параллелизме тензоров и экспертов каждый GPU обслуживает свой шард. Наш декодер разделяет каждую дельту по GPU, которому принадлежит каждое изменение, а каждый GPU проецирует свои изменения прямо в адреса своей упакованной структуры, вплоть до перемешанных (swizzled) масштабных коэффициентов FP8 и FP4.

Устранение конфликтов между синхронизацией дельт весов и обслуживанием

Планировщик SGLang представляет собой цикл событий Python, поэтому любая ресурсоемкая работа на CPU в его процессе конкурирует за GIL и приводит к задержкам планирования батчей. Планирование наших обновлений вместо этого выполняется в отдельном вспомогательном процессе без использования CUDA, который обменивается массивами данных с воркером GPU через общую закрепленную (pinned) память.

Декодирование начинается до завершения загрузки

Дельта поступает в виде множества шардов, и каждый декодируется сразу по мере поступления. Загрузка и декодирование выполняются в конвейере, а не последовательно. Декодер работает в постоянных нативных потоках с низким приоритетом, размер которых подобрался так, чтобы оставить 16 ядер для обслуживания.

Kimi K3 на двух узлах GB300

Kimi K3 занимает 1,7 ТБ в формате развертывания MXFP4. При использовании чекпоинта на CPU каждому узлу GB300 потребуется примерно 838 ГБ хост-памяти только для своей половины весов, что почти исчерпывает весь лимит памяти контейнера обслуживания. Это жесткое ограничение, так как узлы GB300 поставляются с предварительно сконфигурированным аппаратным обеспечением.

Теперь он работает на двух узлах GB300 с четырьмя GPU на каждом. Каждый узел выполняет декодирование, верификацию и подготовку только тех тензоров, которыми владеет, и никогда не хранит веса другого узла. Оба узла должны успешно выполнить запись до того, как новая политика будет опубликована и возобновится обработка запросов.

При дельте в 8,4 ГБ модель Kimi K3 переходит от запроса синхронизации дельта-весов к актуальным весам за 75 секунд при паузе в 4,1 секунды.

Рисунок 5: Каждый узел GB300 декодирует, проверяет и записывает только те тензоры, которыми владеет; ОЗУ хоста не хранит копию весов, а содержит лишь ограниченный буфер дельт и хост-уровень KV.

Выделение памяти под KV-кеш

На протяжении этого блога мы отмечали, что главным преимуществом отказа от хранения чекпоинта в памяти хоста является дополнительный объем памяти, который мы получаем и который позволяет осуществлять выгрузку KV-кеша. Вот как это работает, а также некоторые измерения:

• Вытесненные страницы KV на GPU записываются обратно в память хоста вместо того, чтобы удаляться, и последующий запрос с тем же префиксом перезагружает их вместо повторного запуска префилла.

Вытесненные страницы KV на GPU записываются обратно в память хоста вместо того, чтобы удаляться, и последующий запрос с тем же префиксом перезагружает их вместо повторного запуска префилла.

• Прокатка RL (RL rollouts) представляет собой многошаговые процессы с длинными общими префиксами: перезагрузка вытесненного префикса длиной 150 тыс. токенов из памяти CPU занимает 0,48 секунды по сравнению с 4,9 секунды для свежего префилла.

Прокатка RL (RL rollouts) представляет собой многошаговые процессы с длинными общими префиксами: перезагрузка вытесненного префикса длиной 150 тыс. токенов из памяти CPU занимает 0,48 секунды по сравнению с 4,9 секунды для свежего префилла.

Мы также запустили Kimi K3 на бенчмарке AgentX с выгрузкой на новой архитектуре по сравнению со старой схемой с чекпоинтом на CPU.

• При параллельности 16 KV-кеш уже помещается на GPU, поэтому разницы нет.

При параллельности 16 KV-кеш уже помещается на GPU, поэтому разницы нет.

• При параллельности 48 префиксы агентов больше не помещаются в память GPU, но помещаются в память хоста, поэтому выгрузка превращает промахи в попадания: процент попаданий подскочил со 2% до 92%, а пропускная способность вывода на каждый GPU выросла в 3,2 раза.

При параллельности 48 префиксы агентов больше не помещаются в память GPU, но помещаются в память хоста, поэтому выгрузка превращает промахи в попадания: процент попаданий подскочил со 2% до 92%, а пропускная способность вывода на каждый GPU выросла в 3,2 раза.

Рисунок 6: При параллельности 16 обе конфигурации вмещают всё на GPU; при 48 только реплика с чекпоинтом на GPU имеет свободную память хоста для выгрузки, и попадания в кеш возрастают со 2% до 92%.

Результаты

Все запуски выполняются для одной реплики с тестовыми дельтами в форме тренажера, измерение производится от запроса тренажера на синхронизацию дельта-весов до обслуживания репликой новой политики. Время измеряется от маркера S3 до первого нового ответа.

Если вы хотите внедрить это в свой стек обучения, ознакомьтесь с документацией по прокатке RL или свяжитесь с нашей командой. Если вам кажется интересной такая системная работа, мы нанимаем сотрудников.