AsyncGRPO + LoRA на HF Jobs: обучение без NCCL

Асинхронное обучение с подкреплением (RL) часто упирается в узкое горлышко синхронизации весов между процессами обучения (trainer) и генерации (vLLM). В традиционных кластерах это решается через NCCL, но в среде Hugging Face Jobs (HF Jobs), где каждый контейнер работает на отдельной VM, прямой сетевой связи между нодами нет. Решение от команды Hugging Face — использование LoRA-адаптеров и S3-бакета как общего файлового хранилища. Это позволяет синхронизировать всего несколько мегабайт данных вместо гигабайтов весов модели, исключая необходимость в NCCL. Ниже представлен практический гайд по настройке этой архитектуры, основанный на TRL v1.14. ## Архитектура: Бакет, Прокси и Отсутствие NCCL Ключевая идея заключается в разделении ролей: 1. **Trainer Job**: Запускает `AsyncGRPOTrainer` с LoRA. Периодически сохраняет обновленный адаптер в общую директорию. 2. **vLLM Jobs (Replicas)**: Два или более экземпляров vLLM, которые обслуживают запросы на генерацию (rollouts). Они читают адаптеры из общей директории. 3. **Storage Bucket**: S3-совместимое хранилище, смонтированное в обоих типах Job как FUSE-файловая система. Это заменяет общий сетевой диск (NFS) или InfiniBand. 4. **Proxy Server**: Легковесный прокси, который маршрутизирует запросы на генерацию к реплике, где уже есть соответствующий KV-cache, и транслирует команды загрузки новых адаптеров.

💡

Почему LoRA? Адаптер rank-1 для модели 1.5B весит всего несколько мегабайт. Это позволяет передавать обновления практически мгновенно через файловую систему, в то время как полная модель весит гигабайты, что сделало бы синхронизацию через бакет узким местом.

## Настройка vLLM Replicas Каждая реплика vLLM работает на отдельной GPU. Главная задача — правильно настроить слоты для LoRA-адаптеров. Параметр `max_staleness` в TRL определяет, сколько версий политики (policy versions) может «отставать» от текущей. Если реплика должна удерживать в памяти текущую версию плюс несколько предыдущих, чтобы завершить запущенные rollouts, это влияет на количество необходимых слотов. Формула для `--max-loras` зависит от `max_staleness` и требует резервирования слотов для загрузки новых и выгрузки старых адаптеров. Вот команда для запуска реплики: ```bash # Запуск реплики vLLM с поддержкой runtime LoRA hf jobs run --detach --flavor --timeout --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:latest \ 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 ``` **Важные флаги:** * `VLLM_ALLOW_RUNTIME_LORA_UPDATING=1`: Включает эндпоинт `/v1/load_lora_adapter`. * `VLLM_SERVER_DEV_MODE=1`: Включает эндпоинты `/pause`, `/resume`, `/server_info`, необходимые TRL для управления состоянием. * `--max-loras`: Резервирует слоты для одновременных версий адаптеров. * Версия vLLM должна поддерживать API runtime LoRA.

⚠️

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

## Синхронизация через Storage Bucket Trainer и vLLM используют один и тот же путь в смонтированном бакете. Trainer сохраняет адаптер в директорию `/.vllm_lora/trl-policy-v{N}`, где `{N}` — номер версии. Затем он выполняет атомарное переименование (atomic rename) и отправляет путь в vLLM через эндпоинт `/v1/load_lora_adapter`. Монтирование бакета выглядит так: ```bash # Монтирование бакета в Job hf jobs run ... -v hf://buckets/aminediroHF/asyncgrpo-lora-buckets:/lora ... ``` Все Job должны монтировать бакет в **абсолютно одинаковый путь** (`/lora` в примере). Это позволяет trainer писать файлы, а vLLM читать их без изменения кода TRL или vLLM. Чекпоинты и финальный адаптер также сохраняются в бакет, что позволяет возобновить обучение (resume) после прерывания Job. ## Бенчмарки и Производительность Использование этой архитектуры позволяет значительно ускорить обучение. В тестовом сценарии с моделью Qwen2.5-Math-1.5B и датасетом `sail/Sanity-Test-R1D-1.5B` (1,460 вопросов с уровнем успеха 20-80%): | Метрика | Значение | | :--- | :--- | | **Время обучения (5 запусков)** | От 3 ч 27 мин до 53 мин для 500 шагов | | **Размер синхронизации** | ~несколько МБ (LoRA rank-1) | | **Модель** | Qwen2.5-Math-1.5B | | **GPU** | H200 (на каждую реплику) | | **TRL версия** | v1.14+ | Главный индикатор здоровья системы — метрика `ratio` в TRL. Она должна оставаться близкой к 1. Если она отклоняется, это часто указывает на проблемы с кэшированием KV или рассинхронизацией версий адаптеров. ## Кому подойдёт / что запустится Эта конфигурация идеально подходит для практиков, которые: 1. **Хотят оптимизировать затраты на GPU**: Используя LoRA rank-1, вы можете обучать политики, сопоставимые по качеству с полным fine-tuning, но с меньшими требованиями к памяти и пропускной способности сети. 2. **Работают в среде Hugging Face Jobs**: Архитектура специально спроектирована для ограничения HF Jobs (отсутствие общего локального диска и NCCL между нодами). 3. **Используют модели до 7B**: Для больших моделей размер LoRA-адаптера растет, но все еще остается значительно меньше полного веса. Однако для моделей >7B может потребоваться увеличение `max-loras` и проверка скорости чтения из S3. **Что запустится прямо сейчас:** * Модель: `Qwen/Qwen2.5-Math-1.5B` (или аналогичные малые модели). * Фреймворк: `trl>=1.14`, `vllm` (версия с поддержкой runtime LoRA). * Железо: H200 или A100 (рекомендуется для в

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

AsyncGRPO + LoRA на HF Jobs: обучение без NCCL и за 53 минуты

Читать дальше

Qwen3.8-LiveTranslate: Революция в реальном времени

LLM-as-a-Judge: Автоматическая оценка AI-агентов с Ori Eval

Descript: как 13 моделей работают без инженеров

Главное про AI за 5 минут — раз в день