Назад к блогу

DAPO: как устроена открытая RL-система для обучения больших моделей

DAPO: как устроена открытая RL-система для обучения больших моделей

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

DAPO — это открытая реализация RL-пайплайна для больших языковых моделей, собранная на базе фреймворка verl. В ней четыре ключевые техники — decoupled clip (clip-higher), dynamic sampling, token-level loss и overlong reward shaping — не просто описаны в статье, а доведены до конкретных функций, конфигов и порогов. Ниже разобрана механика: в каком порядке выполняются шаги обучения, как считаются потери и преимущества, как система переживает падения и как воспроизводятся заявленные результаты.

Четыре техники DAPO и где они живут в коде

Все четыре техники включаются через конфигурацию и реализуются в разных местах пайплайна.

Clip-higher задаёт асимметричные границы клиппирования отношения вероятностей: снизу — clip_ratio_low: 0.2, сверху — clip_ratio_high: 0.28. Это позволяет политике сильнее повышать вероятность редких токенов, чем понижать её, и именно эти границы соответствуют ε_low и ε_high в целевой функции DAPO.

Dynamic sampling отбрасывает группы выходов, у которых значения выбранной метрики одинаковы (например, все точности равны 1 или все равны 0), и повторяет генерацию, пока не наберётся достаточно подходящих групп. Включается настройкой algorithm.filter_groups с примером конфигурации в разделе «Dynamic Sampling (with Group Filtering)».

Token-level loss меняет способ сведения матрицы потерь к скаляру: потери усредняются по всем валидным токенам мини-батча, а не по последовательностям. Включается через раздел «Flexible Loss Aggregation Mode (-> Token-level Loss)»: режим агрегации задаётся параметром loss_agg_mode.

Overlong reward shaping штрафует ответы, длина которых превышает допустимую, но ещё укладывается в жёсткий лимит контекста. Работает при включённом overlong_buffer_cfg.enable: тогда вычисляются длина буфера, ожидаемая длина, превышение, коэффициент штрафа и сам штраф, который прибавляется к награде.

Цикл обучения: два варианта порядка шагов

Обучение реализовано в двух трейнерах, и порядок шагов в них различается.

В trainer_base.py цикл step разбит на parameter_sync_step локальных обновлений. Каждое локальное обновление проходит через _step_once, а результаты собираются в KVBatchMeta с идентификатором партиции и наборами ключей и тегов батчей.

Порядок внутри _step_once такой:

  1. сэмплирование батча из replay_buffer (таймер gen);
  2. опционально вычисление награды колокейтед reward-моделью (таймер reward);
  3. балансировка батча по группам data parallel;
  4. вычисление old_log_prob (таймер old_log_prob);
  5. опционально ref_log_prob (таймер ref);
  6. опционально values (таймер values);
  7. вычисление преимуществ advantage (таймер adv);
  8. опционально обновление критика update_critic (таймер update_critic);
  9. обновление актора update_actor (таймер update_actor).

Данные между шагами ходят через TransferQueue. Вычисление награды читает из очереди поля prompts, responses, raw_prompt, а записывает обратно rm_scores и дополнительные ключи награды. Вычисление old_log_prob в bypass-режиме читает rollout_log_probs и записывает old_log_probs, а в decoupled-режиме записывает old_log_probs и entropy.

В ray_trainer.py порядок в fit другой: сначала генерация последовательностей (таймер gen), затем награда (таймер reward), затем либо вычисление old_log_prob (таймер old_log_prob), либо применение bypass-режима, затем ref_log_prob (таймер RefPolicy), затем values (таймер values), затем преимущества (таймер adv), затем update_critic (таймер update_critic) и update_actor (таймер update_actor). Здесь данные объединяются в батч через union: результат генерации, old_log_prob, ref_log_prob и values присоединяются к батчу, а вычисление преимуществ возвращает обновлённый батч.

Восстановление незавершённых промптов после падения

Функция _reissue_inflight_prompts перезапускает промпты, которые остались незавершёнными после восстановления состояния TransferQueue. Сначала она проверяет, что режим не sync, иначе возвращает 0. Затем получает список ключей и выбирает те, у которых тег is_prompt=True и статус pending или running. Для этих ключей она загружает данные, удаляет старые траектории (ключи, у которых is_prompt=False и префикс до _ совпадает с идентификатором незавершённого промпта), после чего заново помещает промпты с тегами is_prompt=True, status='pending' и текущим global_steps. Перед повторной отправкой в батч устанавливается поле global_steps, и запускается генерация последовательностей. Используются поля очереди is_prompt, status и global_steps.

Decoupled clip: асимметричные границы отношения

В compute_policy_loss_vanilla сначала вычисляется отношение вероятностей новой и старой политики как экспонента разности логарифмов. Затем считаются две части цели:

  • неограниченное произведение преимущества на отношение;
  • то же произведение, но с отношением, зажатым снизу через cliprange_low и сверху через cliprange_high.

Итоговая цель берётся как максимум из двух частей. Именно разные границы снизу и сверху и есть decoupled clipping: уменьшение и увеличение вероятности ограничиваются по-разному. Значения по умолчанию — clip_ratio_low: 0.2 и clip_ratio_high: 0.28; нижняя граница отношения вычисляется как 1 - cliprange_low, верхняя — как 1 + cliprange_high.

pg_losses1 = -advantages * ratio
pg_losses2 = -advantages * torch.clamp(ratio, 1 - cliprange_low, 1 + cliprange_high)
pg_losses = torch.maximum(pg_losses1, pg_losses2)

Для отрицательных преимуществ применяется дополнительная нижняя граница clip_ratio_c: считается третий вариант потерь, берётся минимум с уже клиппированным, а выбор между клиппированными вариантами делается по знаку преимущества.

Dynamic sampling: фильтрация групп и добор промптов

При filter_groups.enable=True отбрасываются группы, у которых значения выбранной метрики у всех выходов одинаковы — например, для metric=acc это группы, где все точности равны 1 или все равны 0. Метрика задаётся полем metric; в конфиге перечислены acc, score, seq_reward, seq_final_reward и подобные. При этом reward выбирает каноническую награду до KL-штрафа из rm_scores, а acc, score, seq_reward и прочие — поле из дополнительной информации о награде.

Тренер повторяет сэмплирование с размером gen_batch_size, пока не наберётся достаточно подходящих групп для train_batch_size либо пока не будет достигнут предел max_num_gen_batches. Неположительное значение предела означает отсутствие верхнего ограничения. Если промптов не хватает, счётчик числа генераций увеличивается: пока предел не превышен, генерация продолжается, иначе выбрасывается ошибка о том, что сгенерировано слишком много. Когда промптов достаточно, батч обрезается до train_batch_size * rollout.n.

В примере конфигурации заданы gen_batch_size: 1536, train_batch_size: 512 и filter_groups.metric: acc. Размер gen_batch_size используется как размер батча для загрузчика данных, если он задан, иначе берётся train_batch_size.

Token-level loss: три режима агрегации

Режим loss_agg_mode определяет, как матрица потерь сводится к скаляру.

token-mean усредняет потери по всем валидным токенам мини-батча — через masked_mean, который считает среднее только по позициям с истинной маской. В реализации используется сумма по маске, делённая на число токенов батча и умноженная на размер data parallel группы.

seq-mean-token-sum сначала суммирует потери внутри каждой последовательности с учётом маски, затем берёт среднее по последовательностям. masked_mean здесь не применяется; используется сумма по последовательностям с маской, исключающей полностью замаскированные последовательности.

seq-mean-token-mean внутри каждой последовательности берёт среднее по валидным токенам, затем среднее по последовательностям. Здесь маска последовательности — это число токенов в ней, а потери последовательности делятся на это число с добавкой малой константы для устойчивости.

Разница принципиальна: token-mean даёт каждому токену равный вес, а seq-режимы сначала нормируют внутри последовательности, поэтому длинные ответы не доминируют над короткими.

Overlong reward shaping: линейный штраф за длину

Штраф включается при overlong_buffer.enable: True и наказывает ответы, длина которых превышает допустимую, но ещё укладывается в жёсткий лимит контекста. Сначала берётся длина буфера из конфига, затем вычисляется порог как разность максимальной длины ответа и длины буфера, а превышение — как разность фактической длины и порога. Коэффициент штрафа берётся из конфига, сам штраф считается по формуле и прибавляется к награде:

overlong_buffer_len = self.overlong_buffer_cfg.len
expected_len = self.max_resp_len - overlong_buffer_len
exceed_len = valid_response_length - expected_len
overlong_penalty_factor = self.overlong_buffer_cfg.penalty_factor
overlong_reward = min(-exceed_len / overlong_buffer_len * overlong_penalty_factor, 0)
reward += overlong_reward

Штраф растёт линейно от 0 до overlong_buffer.penalty_factor по мере того, как длина вывода превышает max_response_length - overlong_buffer.len на 0–overlong_buffer.len токенов. В примере конфигурации заданы max_response_length: 20480, len: 4096 и penalty_factor: 1.0.

Отдельно в FAQ отмечено, что большинство экспериментов статьи, включая самый производительный, запускались без Overlong Filtering: он пересекается с Overlong Reward Shaping в части обучения на самых длинных выводах, поэтому его не реализовали.

Оценка преимуществ: реестр и векторизация

Функции оценки преимуществ регистрируются декоратором, который кладёт их в реестр по имени, а достаются они по имени или значению перечисления.

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

compute_grpo_vectorized_outcome_advantage делает то же самое, но векторизованно: получает среднее и стандартное отклонение по группам одним вызовом, вычисляет скаляры по группам и умножает на маску ответа. Обе функции возвращают преимущества дважды — то есть возвращаемые значения равны преимуществам.

KL-контроль: варианты штрафа и трюк с суффиксом

Конфигурация KL-контроля — это dataclass с типом контроля (fixed или adaptive), начальным коэффициентом KL-штрафа, горизонтом для адаптивного контроллера и целевой KL-дивергенцией.

Функция штрафа срезает необязательный суффикс + и передаёт базовое имя дальше. Поддерживаются варианты: kl или k1 (разность логарифмов), abs (абсолютная разность), mse или k2 (половина квадрата разности), low_var_kl или k3 (разность логарифмов с ограничением, экспонента, затем выражение ratio - kl - 1 с ограничением), а full выбрасывает исключение о нереализованности. Выбор варианта определяется тем, что нужно получить: k1 и k3 дают правильное значение KL, но не правильный градиент, тогда как k2 даёт правильный градиент.

Суффикс + меняет поведение: если строка не заканчивается на + или равна mse/k2, возвращается прямой результат; иначе применяется straight-through-трюк, где прямой результат комбинируется с квадратом разности логарифмов. Смысл в том, что математическое ожидание k1 и k3 даёт правильное значение KL, но не правильный градиент, тогда как k2 даёт правильный градиент, — поэтому straight-through связывает k2 с другими методами для несмещённой оценки градиента.

Rollout correction: коррекция off-policy смещения

Конфигурация коррекции обрабатывает смещение из трёх источников: несоответствие rollout-политики и обучаемой политики, устаревание данных из старых чекпоинтов и общий сдвиг распределений.

Уровень агрегации весов задаётся rollout_is: None — только метрики, token — по-токенные веса (низкая дисперсия, смещённые), sequence — по-последовательные (несмещённые, высокая дисперсия). Порог задаётся rollout_is_threshold: одиночное число — это обрезка весов сверху, строка вида lower_upper — обнуление весов вне диапазона. Нормировка весов к среднему 1.0 внутри батча включается флагом rollout_is_batch_normalize. Rejection sampling включается через rollout_rs с порогами rollout_rs_threshold. Режимы различаются по уровню агрегации и по метрике расхождения: token_*, seq_sum_*, seq_mean_*, seq_max_*; суффикс k1 использует расхождение -log r (порог задаётся строкой lower_upper), а k2 и k3 — квадратичное и усечённое расхождения с положительной верхней границей.

Пресеты задают готовые комбинации: decoupled_token_is — по-токенные веса с порогом 2.0 без rejection sampling; decoupled_seq_is — по-последовательные веса с порогом 2.0 без rejection sampling; decoupled_geo_rs — без весов, но с rejection sampling seq_mean_k1 и порогом 0.999_1.001; bypass_ppo_clip — без весов и rejection sampling, полагаясь на PPO-ratio; bypass_pg_is — по-последовательные веса с порогом 2.0 в bypass-режиме с типом потерь reinforce.

В decoupled-режиме коррекция считается один раз на батч по стабильной старой политике. В bypass-режиме этот вызов пропускается, и метрики считает актор, сравнивая текущую политику с rollout-политикой. В bypass-режиме функция потерь читает из конфига тип потерь, уровень весов, пороги, флаг нормировки и режимы rejection sampling. При типе reinforce потери считаются как произведение преимущества, логарифма вероятности и весов; при типе ppo_clip веса передаются как None, потому что отношение текущей политики к старой уже учитывает коррекцию через клиппинг.

Метрики обучения

compute_data_metrics собирает статистики по батчу: суммы по последней оси, а также среднее, максимум и минимум для оценки, награды, преимуществ и возвратов. При этом оценка и награда считаются только по неабортным образцам, а преимущества и возвраты — через выборку по маске ответа. Дополнительно считаются доля абортов и статистики длин ответов неабортных образцов, включая долю клиппированных.

compute_timing_metrics возвращает сырые времена этапов в секундах и время на токен в миллисекундах: для генерации нормировка идёт на число токенов ответа, а для остальных этапов — на общее число токенов (промпт плюс ответ).

compute_throughout_metrics возвращает общее число токенов, время на шаг и пропускную способность как число токенов, делённое на произведение времени и числа GPU.

compute_variance_proxy_metrics вычисляет величину, зависящую от вероятностей токенов старой политики и суммы квадратов вероятностей; при наличии весов коррекции эта величина умножается на квадрат весов. Затем считаются сила сигнала как квадрат нормы градиента и полная мощность: для off-policy как математическое ожидание произведения квадрата весов, квадрата преимущества и величины, для on-policy — без квадрата весов.

MoE load-balance реализован через аккумулятор: обновление накапливает счётчики загрузки из маршрутизированных экспертов и маски ответа, вычисление возвращает метрики по накопленным счётчикам, а добавление метрик записывает число найденных маршрутизированных экспертов и число назначений и сбрасывает аккумулятор.

Функция сведения метрик помечена устаревшей и делегирует в общую утилиту, которая вычисляет среднее каждого списка метрик.

Off-policy staleness и speculative decoding

Метрики устаревания считаются так: сначала определяется, сколько различных версий модели охватывает одна траектория (от минимума до максимума шагов плюс один), затем нижняя граница отставания от текущей политики и худший случай. Для каждого из трёх наборов в метрики добавляются среднее, максимум и минимум.

Метрики speculative decoding считаются только когда доступны данные обо всех трёх величинах — черновиках, принятиях и проверках; если хотя бы одной из них нет, метрики не публикуются. Когда данные есть, по каждому запросу вычисляются доля принятия (ноль при нулевом числе черновиков) и длина принятия (ноль при нулевом числе проверок), после чего публикуются их средние значения. Данные появляются только при включённом MTP и включённом rollout; если дополнительные поля пусты или не все элементы содержат ключ с числом черновых токенов, данные не собираются. Если заполнена лишь часть величин, срабатывает проверка-утверждение.

Как сравнивали и что получилось

В таблице воспроизведения приведены три запуска, различающиеся набором включённых техник, железом, образом окружения и скриптом обучения. Все три используют один и тот же коммит 4f80e4 и переменные окружения из runtime_env.yaml.

ЗапускAIME 2024 Acc.ЖелезоОбразСкрипт
DAPO52%16x8xH800hiyouga/verl:ngc-th2.6.0-cu126-vllm0.8.3-flashinfer0.2.2-cxx11abi0run_dapo_qwen2.5_32b.sh
DAPO w/o Dynamic Sampling50%16x8xH800hiyouga/verl:ngc-th2.6.0-cu126-vllm0.8.3-flashinfer0.2.2-cxx11abi0run_dapo_wo_ds_qwen2.5_32b.sh
DAPO w/o Token-level Loss & Dynamic Sampling44%16x8xH20hiyouga/verl:ngc-th2.5.1-cu120-vllm0.7.4-hotfixrun_dapo_early_qwen2.5_32b.sh

Чтобы повторить замер, нужно перейти в корень репозитория, задать адрес Ray-кластера через RAY_ADDRESS, указать локальную директорию для упаковки через WORKING_DIR, задать runtime-окружение через RUNTIME_ENV и запустить нужный скрипт. RAY_ADDRESS задаёт адрес кластера, к которому выполняется подключение; WORKING_DIR — директорию, которая упаковывается и отправляется на кластер; RUNTIME_ENV — yaml-файл, устанавливающий переменные окружения и pip-пакеты.

cd verl
export RAY_ADDRESS="http://${RAY_IP:-localhost}:8265"
export WORKING_DIR="${PWD}"
export RUNTIME_ENV="./recipe/dapo/runtime_env.yaml"
bash recipe/dapo/run_dapo_qwen2.5_32b.sh

Цифры относятся к конкретным наборам данных: обучающий DAPO-Math-17k и валидационный AIME 2024. В README отдельно подтверждено, что скрипт без Token-level PG Loss и Dynamic Sampling проверен на текущей версии verl и достигает 44 баллов, а полный скрипт DAPO подтверждён на последней версии и даёт 50 баллов.

Подготовка данных и оценка

Подготовка данных выполняется скриптом prepare_dapo_data.sh, который по умолчанию скачивает датасеты в ${HOME}/verl/data. README перечисляет три скрипта: подготовку датасетов, DAPO без Token-level PG Loss и Dynamic Sampling (AIME 44) и полный DAPO (AIME 50).

Для вывода модели используется BytedTsinghua-SIA/DAPO-Qwen-32B, токенизатор загружается из той же модели. Генерация идёт с типом bfloat16, размером тензорного параллелизма 8 и использованием памяти GPU 0.95, а параметры сэмплирования — температура 1.0, top_p 0.7, максимум 20480 токенов.

Для оценки на AIME 2024 модель разворачивается с Ray Serve и vLLM командой запуска сервера с указанием модели и размера тензорного параллелизма 8, а в другом терминале выполняется скрипт оценки с теми же параметрами сэмплирования и файлом теста eval/aime-2024.parquet.

Что из этого следует на практике

  • Порядок шагов различается между трейнерами. В trainer_base.py генерация идёт первой внутри _step_once, а обновление актора откладывается до critic_warmup <= global_steps; в ray_trainer.py порядок тот же по смыслу, но данные объединяются в батч через union, а не через очередь.
  • TransferQueue — это точка восстановления. Незавершённые промпты определяются по тегам is_prompt и статусу, а старые траектории удаляются перед повторной отправкой, поэтому падение не оставляет дублей.
  • Асимметрия клиппирования — это два числа. clip_ratio_low: 0.2 и clip_ratio_high: 0.28 задают разные границы для уменьшения и увеличения вероятности; для отрицательных преимуществ добавляется третья граница clip_ratio_c.
  • Dynamic sampling может не сойтись. Если подходящих групп не хватает, а предел max_num_gen_batches исчерпан, обучение падает с ошибкой, а не продолжается на неполном батче.
  • Режим агрегации потерь меняет вес токенов. token-mean даёт равный вес каждому токену, seq-режимы сначала нормируют внутри последовательности.
  • Штраф за длину линейный и ограниченный. Он растёт от 0 до penalty_factor на интервале в overlong_buffer.len токенов и не применяется за пределами этого интервала.
  • Коррекция off-policy считается по-разному в двух режимах. В decoupled-режиме — один раз на батч по старой политике, в bypass-режиме — актором по текущей политике.
  • Воспроизведение привязано к коммиту и образу. Все три запуска используют коммит 4f80e4, но разные образы и скрипты, а два из трёх — разное железо.

Где смотреть в коде

Источники

Похожее