На пути к обучению с подкреплением (RL) на сэмплерах со смешанной точностью и смешанными ускорителями

В обучении с подкреплением (RL) для больших моделей сэмплеры обычно работают на тех же ускорителях и часто с той же точностью, что и обучающий модуль, что позволяет поддерживать их числовую близость для стабильного обучения. Однако вне сферы RL системы обслуживания не ограничены подобным образом: они могут свободно оптимизировать генерацию по скорости и стоимости, используя форматы низкой точности, такие как FP8, MXFP4 и NVFP4, специально разработанные ядра и даже ускорители, созданные исключительно для префиллинга длинного контекста.

Поскольку сэмплирование уже занимает 80–% графических процессоров в крупных RL-запусках, это оставляет огромный неиспользованный потенциал! Мы задаемся вопросом, можем ли мы использовать сэмплеры со смешанной точностью и смешанными ускорителями, каждый из которых работает на своих собственных ядрах, сохраняя при этом стабильность и производительность обучения.

На пути к этому стоят две проблемы.

Первая проблема — обучение. Каждая из этих оптимизаций увеличивает расхождение между обучением и выводом, делая сэмплы более «off-policy» (вне политики), что часто дестабилизирует процесс обучения. Естественным решением является коррекция off-policy; она позволяет обучающему модулю учиться на данных, сгенерированных другой политикой, подобно устаревшим сэмплам при асинхронном обучении. Поскольку этот метод не делает никаких предположений о том, почему две политики различаются, он должен в равной степени применяться к различным уровням точности или ускорителям.

Вторая проблема — системная. Инфраструктура RL обычно строится вокруг кластеров однородных ускорителей. Смешанные сэмплеры требуют обратного: ускорителей с автоматическим масштабированием, которые при необходимости увеличивают мощность и работают с использованием собственных уровней точности и ядер.

Благодаря коррекции off-policy и нашим бессерверным графическим процессорам с автоматическим масштабированием мы показываем, что стабильное, производительное и более дешевое RL-обучение возможно с использованием сэмплеров со смешанной точностью и смешанными ускорителями, обучая модель с открытыми весами полностью асинхронно на задачах агентного программирования.

Наша установка

Мы обучаем Qwen3.6-35B-A3B на задачах агентного программирования с контекстом до 256 тысяч токенов, используя 2,7 тысячи задач из репозитория Xiaomi released с MiMo-V2.6, после удаления истории репозитория и других следов, через которые некоторые задачи раскрывают эталонное исправление. В каждом эпизоде модель создает патч для проблемы за несколько ходов и отправляет его через bash-инструмент в Modal Sandbox. Она получает вознаграждение 1, если отправленный патч проходит тесты задачи в чистой «песочнице», и 0 в противном случае. Мы проводим оценку на SWE-bench Pro V2, который не пересекается с обучающей выборкой, и отдельно на его сложном подмножестве HARD-51. Каждая контрольная точка оценивается в BF16.

Мы используем Miles в качестве RL-фреймворка с сэмплерами SGLang и обучающим модулем Megatron-LM, который обновляет политику в BF16 на графических процессорах B200. Обучение полностью асинхронное, без ограничений на устаревание данных, и каждое обновление использует 128 промптов по 8 траекторий в каждом. Сэмплеры сохраняют свои собственные ускорители, точность и ядра, а обучающий модуль не повторяет их решения по маршрутизации MoE.

Сэмплеры работают в одной из двух конфигураций:

• Смешанные сэмплеры охватывают шесть типов ускорителей: B200 и B300 обслуживают веса NVFP4, H100 и H200 обслуживают веса FP8 и BF16, а A100 и RTX PRO 6000 обслуживают веса BF16. Каждый сэмплер запускает собственные ядра своего ускорителя и хранит свой KV-кэш в FP8. Маршрутизатор назначает каждую сессию агента одной из этих восьми групп пропорционально пропускной способности группы, поэтому каждая группа обслуживает стабильную долю запросов на протяжении всего обучения.
• Однородные сэмплеры работают только на B200 с весами BF16 и KV-кэшем BF16. Они соответствуют ускорителю и точности обучающего модуля, поэтому их разделяют только ядра и устаревание данных.

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

Обеспечение стабильности RL со смешанными сэмплерами

Мы начинаем с GRPO без какой-либо коррекции расхождения и добавляем три компонента:

• Маскированное importance sampling (MIS) корректирует веса. Оно умножает член каждого токена на его коэффициент важности и маскирует токены, коэффициент которых выходит за пределы фиксированного диапазона, поэтому для токенов, которые остаются, ожидаемое обновление следует политике обучения, а каждый коэффициент остается ограниченным. Без маски несколько токенов, которые сэмплеры генерируют редко, могли бы иметь очень большие коэффициенты и доминировать в обновлении.
• Центрирование оценок (SC) устраняет дрейф. В каждой позиции оно вычитает ожидаемый градиент логарифмической вероятности при политике сэмплирования, поэтому в ожидании на политику влияют только различия в преимуществах между токенами. Коэффициенты важности также устраняют дрейф, но маскирование вновь его вносит.
• Top-p сэмплирование с повтором маски (top-p) исключает редкие токены и дает обеим политикам одинаковую поддержку. Сэмплеры выбирают каждый токен только из набора кандидатов top-p, поэтому редкие токены, которые могут иметь экстремальные коэффициенты, никогда не попадают в данные. Затем обучающий модуль перенормирует свои собственные вероятности по тому же набору. Запуски без этого метода сэмплируют из полного словаря.

Каждый сэмплированный токен добавляет свое преимущество, умноженное на градиент его логарифмической вероятности, к обновлению. В совокупности MIS перевзвешивает этот член коэффициентом важности, SC центрирует градиент, а top-p ограничивает обе политики одними и теми же кандидатами.

Без коррекции обучение разрушается в обеих конфигурациях: при смешанных сэмплерах GRPO разрушается вскоре после начала обучения. При однородных сэмплерах обучение длится намного дольше, прежде чем окончательно разрушиться. Каждый элемент коррекции сам по себе откладывает разрушение, но каждый такой запуск заканчивается на уровне базовой модели или ниже. Только SC + MIS, с top-p или без него, обучаются стабильно до конца: pass@4 растет примерно в течение первых 100–140 шагов, а затем выравнивается на уровне около 65–75% на HARD-51. Полный набор показывает ту же картину с более высокого старта, от 89% до примерно 94%.

Каждое разрушение имеет три общих признака: вознаграждение за обучение падает, KL-дивергенция между сэмплером и обучающим модулем растет в 3 раза или более, и большинство эпизодов заканчиваются повторяющимися ошибками формата. Длина траектории не показывает последовательной закономерности: некоторые разрушенные запуски заканчивают эпизоды раньше, в то время как другие начинают повторяться и длятся дольше. Стабильные запуски не показывают ни одного из этих признаков. Соответствие ускорителю и точности обучающего модуля лишь откладывает разрушение: GRPO на однородных сэмплерах начинается с KL-дивергенции, которая примерно в 4 раза ниже, но следует тем же путем примерно со 100-го шага.

Разрушенный эпизод редко записывает неправильно сформированную команду. Вместо этого в середине задачи модель перестает вызывать свой инструмент, и после трех ответов подряд без вызова инструмента агент завершает эпизод. В контрольной точке GRPO на 40-м шаге при смешанных сэмплерах 62% эпизодов оценки заканчиваются таким образом, в среднем на 22-м ходу, в отличие от 57 ходов для успешных эпизодов. Последний ответ пуст в 98% из них, в то время как только 6% содержат неправильно сформированный вызов.

Наравне с однородными сэмплерами при меньших затратах

Благодаря SC + MIS + top-p смешанные сэмплеры достигают показателей SWE-bench Pro V2, сопоставимых с результатами однородных сэмплеров: значение pass@4 для обоих прогонов остается близким на протяжении всего обучения. Их награда за обучение и норма градиента также следуют схожим траекториям, несмотря на то, что KL между сэмплером и тренером в среднем в 4,7 раза выше. Энтропия совпадает примерно до 160-го шага; после этого она растет у однородных сэмплеров и остается неизменной у смешанных. Важно отметить, что в механизме коррекции нет ничего специфичного для данной модели или этих ускорителей, поэтому мы ожидаем, что эта закономерность будет проявляться более широко.

Отделение сэмплеров от тренера также делает достижение этих результатов менее затратным. На протяжении большей части обучения стоимость обновления на сэмплер была примерно на 15% ниже при использовании смешанных сэмплеров, чем при использовании однородных. Обратите внимание, что наша установка была создана для проверки производительности использования сэмплеров со смешанной точностью и смешанными ускорителями, а не для минимизации затрат, поэтому существуют возможности для дальнейшей экономии: каждый сэмплер можно настроить индивидуально, а состав смеси можно смещать в сторону наиболее экономически эффективных сэмплеров.

В следующих трех разделах объясняются эти результаты: откуда берется рассогласование между обучением и выводом (training–inference mismatch), почему мы не используем чистое on-policy обучение с подкреплением (RL) и как работает off-policy коррекция.

Откуда берется рассогласование между обучением и выводом

Сэмплированный токен может получить одну вероятность от сэмплера и другую от тренера. Причины охватывают всю систему RL:

Смешанные сэмплеры объединяют большинство этих факторов одновременно, и их рассогласование в среднем в 4,7 раза выше, чем у однородных сэмплеров.

Почему не чистое on-policy RL?

Обучение является on-policy, когда каждый токен сэмплируется из обучаемой политики. Когда сэмплеры и тренер представляют собой отдельные системы, для обеспечения on-policy режима требуются две вещи: чтобы сэмплеры вычисляли те же побитовые вероятности, что и тренер (т.е. истинное on-policy RL), и чтобы каждый сэмпл поступал из текущих весов тренера.

Побитово-идентичное сэмплирование и обучение уже были продемонстрированы. Открытые RL-фреймворки делают их идентичными для небольших моделей, а в DeepSeek-V4 были созданы пакетно-инвариантные детерминированные ядра, чтобы сделать обучение и вывод побитово воспроизводимыми. Однако сопоставление каждого источника имеет свою цену:

Большинство этих источников изменяют только вероятности токенов, что учитывается коррекцией. Многоходовой контекст и top-p / top-k сэмплирование нарушают общий контекст и общую поддержку, на которые опирается коррекция, поэтому наши прогоны соответствуют им: токен-в-токен-из, и повтор маски всякий раз, когда мы сэмплируем с top-p. В остальном, сопоставление — это выбор.

На практике крупные прогоны сопоставляют некоторые из этих источников и пытаются корректировать остальные. Прогоны, которые соответствуют точности, маршрутизации и top-p маске, по-прежнему ограничивают устаревание и маскируют токены с экстремальными коэффициентами важности, а прогоны без устаревания по-прежнему корректируют рассогласование между обучением и выводом. В масштабе побитовое сопоставление означает, что ядра пишутся специфично для модели. Но даже в этом случае асинхронный прогон остается off-policy из-за устаревания.

Распространено мнение, что чем ниже это рассогласование, тем стабильнее обучение, поэтому прогоны поддерживают сэмплеры численно близкими к тренеру. Однако стабильность зависит не столько от величины рассогласования, сколько от того, продолжает ли оно расти. Без коррекции рассогласование смещает обновление, что может увеличивать рассогласование до тех пор, пока обучение не разрушится. Меньшее рассогласование лишь откладывает это: GRPO на однородных сэмплерах начинался с KL примерно в 4 раза ниже, чем на смешанных, но все равно позже пришел к краху. С коррекцией обучение на смешанных сэмплерах остается стабильным, а производительность находится на уровне однородных сэмплеров, несмотря на существенное рассогласование.

Величина рассогласования и его рост имеют разные причины. Величина зависит от системы: ускорителя, точности и ядер, на которых работает каждый сэмплер. С другой стороны, рост зависит от обновления, которое и исправляет коррекция. Таким образом, ускорители и точность не требуют нового вида исправления: off-policy коррекция, которую уже применяют крупные прогоны, распространяется и на них. Цель состоит в том, чтобы удерживать рассогласование в ограниченных пределах, а не в том, чтобы поддерживать его низким.

Поэтому мы не стремимся к истинному on-policy RL: сэмплеры и тренер созданы для разных задач, и каждый продолжает совершенствоваться в своей. Сопоставление связывает их вместе, и каждый отказывается от части своих оптимизаций; с коррекцией каждый сохраняет их.

Как работает off-policy коррекция

Обучение на данных, сгенерированных другой политикой, — это устоявшаяся проблема в RL. IMPALA отделили акторов, генерирующих опыт, от обучающегося и корректировали отставание политики, вместо того чтобы держать их в жесткой связке. Политика сэмплирования q генерирует токены, а политика обучения pθ учится на них. При смешанных сэмплерах каждый токен поступает от одного сэмплера с одной версией весов, где каждый сэмплер записывает вероятность токена, поэтому q может отличаться от токена к токену. Коррекции никогда не нужно знать, почему q отличается от pθ: ядра, точность, ускорители и устаревание — все это проявляется только как другой q.

Когда токены поступают из q вместо pθ, обновление градиента политики в каждой позиции приводит к двум ошибкам.

Дрейф. В рамках политики обучения градиент логарифмической вероятности в среднем равен нулю в каждой позиции; в рамках q он в среднем направлен в сторону q. Таким образом, часть обновления перемещает политику обучения независимо от того, какие токены получили награду: в сторону политики сэмплирования, когда ожидаемое преимущество положительно, и в сторону от нее, когда оно отрицательно. Нормализация преимуществ внутри группы не устраняет это, поскольку ожидаемое преимущество после заданного префикса, как правило, не равно нулю. Затем сэмплеры загружают новые веса тренера и снова добавляют свою ошибку, поэтому тренер продолжает двигаться к цели, которая движется вместе с ним.

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

Обе ошибки исчезают, когда q = pθ, и растут в противном случае.

Importance sampling исправляет обе ошибки путем взвешивания члена каждого токена его отношением pθ / q, при условии, что сэмплеры могут сэмплировать каждый токен, который может сэмплировать политика обучения. Взвешивание всей траектории произведением ее отношений является точным, но его дисперсия может расти экспоненциально с длиной, поэтому крупные прогоны применяют отношение по одному токену за раз, аппроксимация первого порядка, ошибка которой растет вместе с рассогласованием и длиной траектории. Даже для каждого токена, токен, который сэмплеры производят редко, может иметь вес до 1 / q. Ограничение веса путем его усечения или маскирования контролирует дисперсию, но вновь вводит дрейф. MIS маскирует: он удерживает отношение в заданном диапазоне и устанавливает его в ноль за его пределами.

Центрирование оценок (score centering) вместо этого устраняет дрейф. В каждой позиции оно вычитает ожидаемый градиент логарифмической вероятности при q, поэтому в ожидании постоянное преимущество больше не смещает стратегию. Оно сохраняет неверное взвешивание, поэтому его используют в паре с MIS: центрирование взвешенного члена устраняет дрейф независимо от веса, включая дрейф, который возвращает маскирование. SC + MIS устраняет дрейф и взвешивает сохраняемые токены в соответствии с обучающей стратегией. Вычисление ожидаемого градиента при q требует наличия q для каждого токена. Сэмплеры записывают его для своих наиболее вероятных токенов; для всех остальных токенов q аппроксимируется вероятностями тренера, пересчитанными относительно оставшейся вероятностной массы сэмплеров.

Top-p сэмплирование нарушает общую поддержку, а повторное воспроизведение маски (mask replay) восстанавливает её. Сэмплеры выбирают токены только из top-p набора, поэтому токены вне его имеют нулевую вероятность при q. Повторное воспроизведение маски в тренере перенормирует pθ по тому же набору. Самые редкие токены, которые могут иметь наибольшие коэффициенты, никогда не попадают в данные, а сэмплеры записывают вероятность каждого кандидата, поэтому центрирование является точным, а не оценочным.

Что показывают наши запуски

• GRPO сохраняет обе ошибки. Норма его градиента возрастает примерно в 20 раз вскоре после начала обучения.
• GRPO + MIS исправляет взвешивание, но не дрейф, который возвращает его маска. Норма его градиента остается в пределах нескольких значений от начальной, но в процессе обучения энтропия падает, так как KL-дивергенция и доля токенов с коэффициентом вне диапазона [0.2, 5] возрастают более чем в 10 раз.
• SC устраняет дрейф, но сохраняет неверное взвешивание. Его энтропия неуклонно падает, а KL-дивергенция возрастает примерно в 3 раза перед тем, как запуск терпит крах.
• GRPO + top-p исключает самые редкие токены. В то время как норма градиента GRPO взлетает, его собственная остается в 3–10 раз меньше, но KL-дивергенция неуклонно растет, и он терпит крах позже.
• SC + MIS устраняет дрейф и исправляет взвешивание для сохраняемых токенов. Он остается стабильным, за исключением одного кратковременного скачка градиента, в то время как его KL-дивергенция возрастает примерно в 2 раза.
• SC + MIS + top-p также выполняет точное центрирование. У него самая низкая KL-дивергенция, а норма градиента и энтропия остаются неизменными.

SC + MIS + top-p обучается при среднем рассогласовании в 4.7 раза большем, чем у однородных сэмплеров, и остается стабильным. В каждом запуске, который терпит крах, рассогласование возрастает в 3 раза или более.

Для каждой группы сэмплеров работа по исправлению следует за рассогласованием. Доля группы в токенах с экстремальными коэффициентами больше зависит от ее KL-дивергенции, чем от доли в выборках: H100 FP8 дает 10% выборок, но 29% таких токенов, в то время как B200 и B300 NVFP4 дают 47% выборок и 16%. Даже для H100 FP8 только около 0.3% его токенов выходят за пределы [0.2, 5]. Исправление взвешивает каждый токен по его собственному коэффициенту, поэтому одно и то же исправление сохраняет стабильность обучения при любом составе групп.

Ограничения

MIS, SC и top-p работают по одному токену за раз, поэтому каждый из них наследует ограничение поаппроксимации для каждого токена: префиксы и преимущества по-прежнему поступают от сэмплеров. Исправления на уровне последовательностей, такие как GSPO и , а также ограничение устаревания защищают от этой ошибки, когда она становится слишком большой. В наших запусках ни одно из них не используется, и они остаются стабильными при контекстах до 256 тысяч токенов и среднем устаревании около 4.

Top-p сэмплирование с повторным воспроизведением маски также меняет то, что обучается: обновление оптимизирует стратегию, усеченную до top-p набора кандидатов, а логиты вне его не получают градиента в этой позиции. Наша оценка проводится с использованием выборок из полного словаря, и SC + MIS + top-p показал результаты примерно такие же, как SC + MIS, который также использует выборки из полного словаря при обучении.

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

Эти исправления — не единственные варианты, но их объединяет один принцип: рассогласование может быть большим, пока обновление остается близким к on-policy, а рассогласование остается ограниченным. Вероятно, появятся и более совершенные методы.

Что дальше

Исправление отделяет сэмплеры от тренера, и это меняет то, на чем может выполняться обучение с подкреплением (RL). Сэмплирование больше не требует большого однородного кластера: любой ускоритель, способный обслуживать модель, может предоставлять выборки, поэтому парк сэмплеров может быть масштабирован отдельно от тренера и увеличен в процессе запуска по мере необходимости (например, по мере удлинения траекторий). Оптимизации вывода также напрямую переносятся в RL: более быстрые ядра, спекулятивное декодирование, более низкая точность или разделение префилла и декодирования между разными ускорителями могут ускорить сэмплирование без необходимости подстраиваться под них со стороны тренера. А тренер может перейти на лучшие ускорители и собственную низкую точность, такую как FP8 или NVFP4, без необходимости следовать этому для сэмплеров. Поскольку ускорители продолжают диверсифицироваться по поставщикам, поколениям и форматам чисел, это разделение становится необходимым для обеспечения обучения передовых моделей для всех.

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

Благодарности

Мы хотели бы поблагодарить команду Miles, Люка Хуана, Мартина Марека, Дхрува Гаутама, Сюю Ли, Байрона Сю, Цяна Ли и Цилинь Чжу за обсуждение черновика этой статьи.

Ссылки: nearly identical for GLM-5