Асинхронный GRPO с LoRA на Hugging Face Jobs: настройка vLLM, прокси и Storage Buckets

В релизе библиотеки TRL v1.14 (PR #7017) появилось важное обновление для асинхронного оптимизатора AsyncGRPOTrainer — полноценная поддержка адаптеров LoRA. Теперь система способна обучать компактный LoRA-модуль вместо полной весовой матрицы языковой модели и передавать на инференс-серверы vLLM исключительно веса адаптера.

Использование LoRA в задачах RL (обучения с подкреплением) обосновано исследованиями: например, работы лаборатории Thinking Machines («LoRA Without Regret») показывают, что в алгоритмах типа policy-gradient адаптеры низкого ранга (даже r=1) демонстрируют результаты, не уступающие полному файнтюнингу. Это объясняется тем, что функция преимущества (advantage function) транслирует ограниченный объём информации за эпизод (порядка O(1) бит), и ёмкости ранг-1 адаптера вполне достаточно для её усвоения.

С инженерной точки зрения это дает колоссальное преимущество. Если полная модель на 1.5 млрд параметров занимает в памяти порядка 3 ГБ, то её LoRA-адаптер 1-го ранга весит всего несколько мегабайт. Вместо прокачивания гигабайтов весов по сети после каждого оптимизационного шага нам достаточно передать tiny-адаптер. При этом vLLM умеет удерживать несколько адаптеров в памяти одновременно, позволяя незавершённым роллаутам дорабатывать на той версии политики, с которой они начинались.

В стандартных локальных кластерах AsyncGRPOTrainer связывает процесс обучения и генерации через общую файловую систему или NCCL-группу. Однако при использовании облачной инфраструктуры Hugging Face Jobs каждая задача запускается в изолированном контейнере на отдельной виртуальной машине без прямого сетевого взаимодействия NCCL между узлами. Благодаря передаче только LoRA-адаптеров нам больше не нужен межсерверный NCCL: достаточно примонтировать сетевой бакет (Storage Bucket) во все контейнеры как файловую систему POSIX. Ниже я подробно разберу, как устроена эта архитектура и как нам удалось разогнать её в 3.9 раза.

Архитектура: связь Hugging Face Jobs и Storage Buckets

В TRL новый механизм синхронизации передаёт на vLLM не тензоры, а пути к файлам. Через заданное число шагов оптимизатора тренер сохраняет адаптер в директорию <output_dir>/.vllm_lora/trl-policy-v{N}, атомарно переименовывает её и отправляет HTTP-запрос на эндпоинт vLLM /v1/load_lora_adapter. Сервер vLLM считывает адаптер прямо с диска, после чего генерация начинает использовать имя model="trl-policy-v{N}".

На Hugging Face Jobs общая файловая система организуется путем монтирования бакета через утилиту hf-mount. Каждый контейнер получает доступ к одной и той же папке:

# Монтирование общего бакета во все задачи HF Jobs
hf jobs run ... -v hf://buckets/my-org/asyncgrpo-lora-buckets:/lora ...

Благодаря этому тренер пишет в /lora/run_name/.vllm_lora/, а инференс-серверы мгновенно считывают данные по тому же пути. Кроме того, в бакете автоматически сохраняются чекпоинты: если обучающая задача прервётся, процесс подхватит актуальный адаптер из бакета без потери прогресса.

Конфигурация трех рабочих задач HF Jobs

Настройка инстансов vLLM

Для генерации используются изолированные GPU-ноды с официальным образом vllm/vllm-openai:v0.27.1. Нам необходимо активировать динамическую загрузку LoRA и выделить нужное количество слотов под адаптеры.

Количество слотов рассчитывается исходя из параметра max_staleness. Если max_staleness=4, то роллауты, созданные на версии v3, остаются валидными для обучения, даже когда тренер ушёл вперёд до версии v7. Чтобы vLLM не выгружал активные веса, требуется удерживать в памяти текущую версию плюс 4 предыдущие, а также 1 дополнительный слот для выполнения атомарной смены адаптеров во время загрузки. Итого задаём --max-loras 6.

# Команда запуска реплик vLLM на Hugging Face Jobs
hf jobs run --detach --flavor h200 --secrets HF_TOKEN \
  --expose 8000 \
  -v "hf://buckets/${BUCKET}:/lora:ro" \
  -e VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 \
  -e VLLM_SERVER_DEV_MODE=1 \
  -- vllm/vllm-openai:v0.27.1 \
  vllm serve Qwen/Qwen2.5-Math-1.5B --host 0.0.0.0 --port 8000 \
  --max-model-len 4096 --logprobs-mode processed_logprobs --generation-config vllm \
  --enable-lora --max-lora-rank 1 --max-loras 6

Обратите внимание: фиксировать имя адаптера и перезаписывать веса под одним постоянным наименованием нельзя. Дело в том, что vLLM привязывает кэш префиксов (KV cache) к имени модели/адаптера. Если имя не меняется, блоки KV, рассчитанные на старых весах, будут ошибочно использоваться для новых весов, из-за чего возникнет дрейф коэффициента важности (ratio) и нарушение математики GRPO. Уникальные версионированные имена trl-policy-v{N} полностью исключают эту проблему.

Выбор датасета: датасет Sanity

В качестве тестового набора выбран датасет sail/Sanity-Test-R1D-1.5B (из работы Qi et al., 2025, «Precision-RL»). В нём собраны 1 460 математических задач из MATH, для которых модель DeepSeek-R1-Distill-Qwen-1.5B показала процент успешных ответов от 20% до 80%. Этот выбор идеален для валидации RL: задачи не являются слишком простыми или безнадёжно сложными, что обеспечит чёткий градиентный сигнал.

Параметры эксперимента: модель Qwen/Qwen2.5-Math-1.5B, LoRA ранг 1, alpha 2, скорость обучения (learning rate) 4e-5, 8 генераций на промпт, 128 комплишенов за шаг, максимум 3 000 сгенерированных токенов при контексте 4 096 токенов.

Настройка модуля обучения

Обучающий скрипт использует базовый AsyncGRPOTrainer из библиотеки TRL:

from peft import LoraConfig
from trl.experimental.async_grpo import AsyncGRPOConfig, AsyncGRPOTrainer

config = AsyncGRPOConfig(
    output_dir="/lora/sanity-lora-r1",
    vllm_server_base_url="http://localhost:8000", # Обращение к локальному прокси
    max_staleness=4,
    weight_sync_steps=4,
    save_strategy="steps",
    save_steps=50,
)

trainer = AsyncGRPOTrainer(
    model="Qwen/Qwen2.5-Math-1.5B",
    args=config,
    peft_config=LoraConfig(r=1, lora_alpha=2, target_modules="all-linear"),
)

Кастомный прокси-сервер маршрутизации

Поскольку тренер и инференс-серверы находятся на разных машинах HF Jobs, напрямую связать их стандартным образом нельзя по двум причинам:

1. Открытые порты HF Jobs требуют передачи заголовка Authorization: Bearer <HF token> в каждом HTTP-запросе.
2. Параллелизм генерации: в режиме data_parallel_size > 1 запрос /v1/load_lora_adapter приходит только на один DP-ранк, а остальные остаются на старых весах. Кастомный прокси решает эту проблему, рассылая адаптер сразу на все реплики.

Прокси запускается локально на хосте тренера (127.0.0.1:8000). Для TRL он выглядит как одиночный vLLM-сервер, но внутри себя он выполняет две критические функции: умную маршрутизацию запросов и вещание команд синхронизации.

Маршрутизация роллаутов по KV-префиксу

Процесс генерации текста состоит из двух этапов: префилла (prefill) и декодинга (decode). На этапе префилла модель параллельно просчитывает ключи и значения (KV cache) для всех токенов промпта. Если две задачи содержат одинаковый начальный фрагмент текста (например, системный промпт или шаблон чата), реплика vLLM, у которой этот префикс уже находится в памяти, может полностью пропустить этап префилла.

Поскольку алгоритм GRPO генерирует пачку из G ответов (в нашем случае G=8) на один и тот же вопрос, важно направить все 8 запросов на одну и ту же реплику. В противном случае ресурсы GPU будут тратиться впустую на повторный просчёт идентичного префилла.

Маршрутизатор работает по следующему алгоритму:

1. Разбиение на блоки: текст промпта режется на блоки по 16 токенов (аналогично внутреннему устройству кэша vLLM).
2. Хэширование префикса: последовательно рассчитываются сцепленные хэши блоков. Начальным значением (seed) служит имя текущего LoRA-адаптера. Это гарантирует, что кэш от версии v3 никогда не пересечётся с версией v4.
3. Сравнение промптов: первые токены (системный промпт, chat template) совпадают у большинства задач в датасете.
4. Фильтрация общего префикса (Common Prefix): блоки, присутствующие на всех репликах или имеющие разветвления continuation, признаются общими и исключаются из логики привязки, так как они не уникальны для конкретного вопроса.
5. Выбор реплики:
   • Affinity hit: если реплика уже имеет уникальные блоки промпта в кэше и ее очередь задач не перегружена (отставание не более 8 запросов от наименее загруженной), запрос отправляется на неё.
   • Spill: если целевая реплика перегружена, кэшем жертвуют и запрос сбрасывается на свободный сервер.
   • Unmatched: для нового вопроса запрос отправляется на наименее загруженную реплику (по принципу round-robin).
6. Повторное использование префилла: первый запрос из группы просчитывает префикс, а оставшиеся 7 моментально переходят к стадии декодинга.

# Упрощенный фрагмент логики выбора реплики в прокси-сервере
def choose(self, upstreams, model, prompt):
    hashes = self.block_hashes(model, prompt)
    matched = self.matched_prefix(hashes)
    common = self.common_prefix_len(hashes)
    specific = [max(0, m - common) for m in matched]
    
    least = min(u.inflight for u in upstreams)
    best = max(range(self.n), key=lambda i: (specific[i], -upstreams[i].inflight))
    
    if specific[best] > 0 and upstreams[best].inflight - least <= self.cfg.imbalance:
        pick = best # Попадание в кэш (Affinity hit)
    else:
        candidates = [i for i in range(self.n) if upstreams[i].inflight == least]
        pick = candidates[self.rr % len(candidates)] # Сброс перегрузки или новый промпт
        self.rr += 1
        
    return upstreams[pick]

Вещание обновлений адаптера

Операция загрузки адаптера выполняется по принципу «всё или ничего». Если при вызове /v1/load_lora_adapter хотя бы одна реплика возвращает ошибку (например, бакет ещё не успел синхронизировать файлы на данной ноде), прокси делает повторные попытки. В случае неустранимого сбоя прокси автоматически выгружает адаптер со всех остальных реплик, предотвращая рассинхрон политики.

Запросы /pause, /resume и /v1/unload_lora_adapter также вещаются на все узлы, а эндпоинт /health выдаёт статус 200 только при исправности каждой реплики.

Результаты полных прогонов и метрики

Синхронизация весов и эффективность маршрутизатора

За 500 шагов обучения было выполнено 126 операций синхронизации адаптеров (всего 252 успешные загрузки на 2 репликах vLLM). В среднем задержка полного цикла синхронизации составила 8.5 секунд.

Статистика маршрутизатора за 64 728 сгенерированных роллаутов:

• Попадания в кэш (Affinity hits): 84.5%
• Новые промпты (Unmatched): 14.2% (при теоретическом минимуме 12.5% для пачек из 8 запросов)
• Сброс из-за перегрузки (Spills): 1.3%

Награда и проверка корректности

За 500 шагов средняя награда модели выросла с 0.145 до 0.438. Главный показатель корректности асинхронного GRPO — значение ratio (отношение вероятностей токенов в новой и старой политиках). На протяжении всех 500 шагов коэффициент находился в диапазоне от 0.9993 до 1.0004, удерживаясь строго вблизи идеальной единицы. Средняя задержка политики (staleness) составила 1.5 версии при предельно допустимых 4.

Оптимизация узких мест: от 3.5 часов до 53 минут

Асинхронный RL — это конвейер между тренером и генератором. Для поиска узких мест используются метрики perf/rollout_wait_s (время ожидания данных тренером) и rollout/backpressure_s (время простоя генератора из-за переполнения очереди). В ходе серии из 5 экспериментов мы последовательно устранили узкие места системы.

Запуск 1 (r1-dp2) — медленный тренер

В базовом конфиге время шага составляло 22.9 с, из которых 21.9 с уходило на прямой и обратный проход (forward/backward). Очередь роллаутов была заполнена (476 из 512), а генерация простаивала по 15 секунд под противодавлением (backpressure). Полезная утилизация GPU тренера (MFU) составила катастрофические 3.9% из-за обработки микробатчей размера 1 (per_device_train_batch_size=1, 64 микробатча за шаг).

Запуск 2 (r1-dp2-tb16k) — упаковывание микробатчей

Мы включили упаковку токенов (token-budget batching): token_budget=16384 и gradient_accumulation_steps=6. Вместо 64 микробатчей по 1 последовательности тренер стал обрабатывать 6 плотно упакованных рядов (~12.7 последовательностей на ряд). Время forward+backward упало с 21.9 с до 5.6 с, MFU вырос до 19%, а скорость генерации на vLLM подскочила с 4.6k до 25k токенов/сек просто за счёт освобождения очереди!

Запуск 3 (r1-dp2-tb16k-nockpt) — отключение чекпоинтинга градиентов

Соотношение времени backward к forward составляло 3.2 вместо нормальных 1.0–2.0. Причина кроется в дефолтном параметре gradient_checkpointing=True, который повторно пересчитывал forward проход. Так как модель 1.5B легко помещается в память H200 (25 ГБ из 141 ГБ), мы отключили чекпоинтинг (gradient_checkpointing=False). Время forward+backward сократилось до 4.6 с, а MFU поднялся до 23%. Буфер опустел, и узкое место сместилось на сторону генерации vLLM.

Запуск 4 (r1-dp3-tb16k-nockpt) — добавление третьей реплики vLLM и скрытое ограничение

Мы добавили третью ноду vLLM, однако скорость генерации практически не изменилась (выросла с 25k до 26k ток/сек). Посмотрев на метрику rollout/inflight, мы обнаружили жесткий лимит 128 параллельных запросов, который разделился по ~43 запроса на каждую из трёх реплик. Узким местом оказался консервативный клиентский лимит в коде тренера.

Запуск 5 (r1-dp3-inflight384) — снятие лимита параллельных запросов

Мы увеличили max_inflight_tasks=384 и queue_maxsize=768. Каждая из трех реплик vLLM получила по 128 параллельных задач. Очередь моментально заполнилась до ~690 элементов, время шага оптимизатора снизилось до 4.8 секунд, а общее время выполнения 500 шагов сократилось с 3 часов 27 минут до 53 минут!

Сводная таблица результатов и инструкции по запуску

Для воспроизведения оптимизированного запуска на Hugging Face Jobs используйте следующие команды:

# Клонирование репозитория и авторизация
git clone https://github.com/AmineDiro/hfjobs-lora-buckets && cd hfjobs-lora-buckets
hf auth login

# Запуск финального оптимизированного конфига (Запуск 5, ~53 минуты)
TOKEN_BUDGET=16384 GRAD_ACCUM=6 GRADIENT_CHECKPOINTING=0 PROXY_LORA_RETRY_S=0.5 \
MAX_INFLIGHT=384 QUEUE_MAXSIZE=768 MAX_STEPS=500 ./run_all.sh -- wait

Связка LoRA-адаптеров, монтируемых сетевых бакетов и умного HTTP-прокси позволяет эффективно масштабировать асинхронное RL-обучение языковых моделей без использования дорогостоящих межсерверных NCCL-соединений.

Источник: huggingface.co