AsyncGRPO с LoRA на HF Jobs: ускорение обучения в 4 раза без NCCL

Суть прорыва: разделение обучения и инференса

В версии TRL v1.14 появилась поддержка AsyncGRPOTrainer с адаптерами LoRA. Ключевое отличие от классических подходов — синхронизация весов происходит не через высокоскоростную сеть NCCL, а через файловую систему. Адаптер ранга-1 для модели 1.5B весит всего несколько мегабайт (против 3 ГБ полной модели), что позволяет передавать его через Storage Bucket Hugging Face Jobs, смонтированный как FUSE-файловая система.

Это решает проблему изоляции Hugging Face Jobs: каждый Job работает на отдельной VM, и у них нет общего локального диска или возможности создать NCCL-группу. Архитектура использует бакет как общий раздел, где тренер пишет адаптер, а серверы vLLM читают его.

Архитектура: Бакет, Прокси и vLLM

Решение состоит из трех компонентов, работающих асинхронно:

• Trainer Job: Запускает AsyncGRPOTrainer (с FSDP). Каждые несколько шагов оптимизатора он сохраняет адаптер в директорию .vllm_lora/trl-policy-v{N} и делает атомарный редирект пути.
• vLLM Jobs (Replicas): Два экземпляра vLLM (v0.27.1) на GPU H200. Они монтируют тот же бакет в режиме чтения. Поддерживается VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 для загрузки адаптеров на лету.
• Proxy Server: Маленький прокси, который маршрутизирует запросы (rollouts) к реплике, где уже есть нужный KV-prefix, и рассылает обновления адаптеров всем репликам.

Почему это важно для RL

Исследования (например, от Thinking Machines) показывают, что для policy-gradient RL адаптера ранга-1 достаточно, так как функция преимущества дает лишь ~O(1) бит информации на эпизод. Это позволяет использовать крошечные веса для обучения, избегая передачи гигабайтов данных между узлами.

Результаты производительности

Использование LoRA и асинхронной синхронизации через бакет дало кратный прирост скорости. В тестовых запусках (5 runs) один и тот же рецепт обучения на 500 шагов сократился с 3 часов 27 минут до 53 минут.

Технические детали настройки

Для корректной работы необходимо учитывать max_staleness. Если установлено значение 4, vLLM должен держать в памяти текущую версию политики плюс 4 предыдущие, чтобы завершить запущенные rollouts. Это требует --max-loras 6 (4 старых + 1 новая + 1 слот для переключения). Команда запуска vLLM выглядит следующим образом:

hf jobs run --detach --flavor h200 --timeout 8h \
    -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 \
        --enable-lora --max-lora-rank 1 --max-loras 6

Такой подход позволяет масштабировать RL-обучение на эфемерных инфраструктурах, таких как Hugging Face Jobs, без необходимости в сложных сетевых настройках между узлами.

Источник: Hugging Face blog ↗