Что изменилось в AsyncGRPOTrainer и LoRA
Недавнее обновление библиотеки TRL (Transformer Reinforcement Learning), версии 1.14, добавило поддержку обучения LoRA-адаптеров через AsyncGRPOTrainer. Объясним, что это означает. AsyncGRPOTrainer — это инструмент для обучения языковых моделей с использованием асинхронного градиентного метода. LoRA (Low-Rank Adaptation) — это способ обучать небольшие адаптеры внутри большой модели вместо полного переобучения всей модели. Такой адаптер — по размеру несколько мегабайт, а сама модель — несколько гигабайт.
Из-за этого теперь AsyncGRPOTrainer обучает только адаптер, а не всю модель, и для передачи обновлений на сервера вывода (vLLM) отправляет только адаптер. Это решает ранее существующую проблему передачи больших объемов данных и синхронизации между машинами.
Разделение тренировки и вывода на разные машины в Hugging Face Jobs
Hugging Face Jobs — это технология, позволяющая запускать контейнеры с ИИ-задачами на отдельных виртуальных машинах (ВМ). Особенность — одна ВМ = один Job, без возможности «разделять» работу внутри Job по нескольким узлам или машинам.
Раньше AsyncGRPOTrainer и сервера вывода vLLM запускались вместе на одном сервере или в рамках одного кластера с общими дисками или средствами передачи данных (например, NCCL). Теперь же делают иначе: – Тренировочный Job обучает LoRA-адаптер. – Выводные Jobs — два сервера vLLM — на других машинах держат базовую модель и адаптеры. – Между ними нет прямой связи через NCCL (коммуникационную библиотеку для GPU-кластеров). – Для синхронизации адаптеров используется технология Storage Buckets — удаленный облачный диск, смонтированный в каждый Job, как обычная папка.
Этот подход позволяет тренировать и генерировать ответы на разных машинах, не передавая по сети гигабайты данных модели, а только мегабайты адаптеров.
Роль прокси и как синхронизируются данные
Чтобы гарантировать корректную работу системы, поставили прокси-сервер: – Он добавляет заголовки аутентификации для запросов к серверам вывода; – Направляет запросы на ту копию модели vLLM, которая уже хранит нужный промежуточный ключевой контекст (KV-префикс), чтобы избежать вычислений заново; – Рассылает обновления адаптеров от тренера всем серверам вывода.
AsyncGRPOTrainer сохраняет адаптер как файлы в общем Storage Bucket, а сервера вывода по специальному API подгружают свежие версии адаптеров с диска. Это гарантирует, что при каждом обновлении тренера все vLLM реплики получают актуальные веса.
Почему LoRA подходит для реинфорсмент-обучения и какие выгоды
LoRA хорошо работает для RL (обучения с подкреплением), потому что адаптер маленький и обучается быстро. Исследование от Thinking Machines показало, что rank-1 адаптер (очень низкая размерность) может соперничать с полным тонким тюнингом в задачах с градиентом политики.
По сути, информация в сигнале обучения RL ограничена — достаточно небольшого объёма, чтобы эффективно улучшать модель. Поэтому обучать малый адаптер разумнее, либо даже быстрее и с меньшим расходом ресурсов.
Кроме того, такой подход экономит время и трафик при синхронизации между машинами: если передавать всю модель — это гигабайты за обновление, а с LoRA — всего мегабайты.
Как попробовать такую систему и где она пригодится
Если вы работаете с Reinforcement Learning и используете TRL, можно перейти с полного тонкого тюнинга на LoRA-адаптеры.
– Установите версию TRL не ниже 1.14; – Запустите AsyncGRPOTrainer с опцией LoRA в рамках Hugging Face Jobs; – Подключите в каждый Job общий Storage Bucket для обмена адаптером; – Настройте прокси-сервер для маршрутизации запросов и передачи обновлений; – Используйте два сервера vLLM для масштабирования вывода.
Это хорошо подойдет для задач, где тренировка и генерация разделены: например, в продакшн-системах с ограниченными ресурсами на каждом сервере, с желанием разгрузить inference и повысить пропускную способность.
Кому это актуально сейчас и что делать дальше
Подход актуален для инженеров и исследователей, которые обучают RL-политики на крупных языковых моделях и хотят оптимизировать инфраструктуру под облачные среды с ограничениями масштабирования на узлы.
Если вы просто пользуетесь готовыми чатами и не занимаетесь обучением, эта тема пока не для вас — но полезно знать, как работают современные распределённые системы обучения и вывода.
Для желающих вникать глубже полезно познакомиться с основами LoRA, ассинхронного обучения, а также работой с Hugging Face Jobs и их Storage Buckets.
Источник: https://huggingface.co/blog/asyncgrpo-lora-hfjobs, Amine Dirhoussi и др.