TRL documentation
Ускорение обучения
Ускорение обучения
В этом руководстве описаны различные методы ускорения обучения в TRL. Каждая техника сопровождается минимальными примерами и ссылками на более подробную документацию.
vLLM для быстрой генерации в онлайн-методах
Онлайн-методы, такие как GRPO или Online DPO, требуют от модели генерации завершений, что часто является медленным процессом и может существенно влиять на время обучения. Чтобы ускорить генерацию, вы можете использовать vLLM — библиотеку, которая обеспечивает быструю генерацию, в том числе благодаря PagedAttention. Онлайн-тренажеры TRL поддерживают vLLM, что значительно повышает скорость обучения. Подробнее см. в разделе vLLM Integration.
Чтобы использовать vLLM, сначала установите его с помощью:
pip install trl[vllm]
Сначала запустите сервер vLLM, выполнив следующую команду:
VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \
--weight-transfer-config '{"backend": "nccl"}' \
--logprobs-mode processed_logprobs \
--max-logprobs -1Затем запустите скрипт обучения и передайте use_vllm=True в аргументы обучения.
from trl.experimental.online_dpo import OnlineDPOConfig
training_args = OnlineDPOConfig(..., use_vllm=True, vllm_mode="server")Сначала запустите сервер vLLM, выполнив следующую команду:
VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \
--weight-transfer-config '{"backend": "nccl"}' \
--logprobs-mode processed_logprobs \
--max-logprobs -1Затем запустите скрипт обучения и передайте use_vllm=True в аргументы обучения.
from trl import GRPOConfig
training_args = GRPOConfig(..., use_vllm=True, vllm_mode="server")Вы можете настроить конфигурацию сервера, передав дополнительные аргументы. Подробнее см. в разделе интеграция с vLLM.
При использовании vLLM убедитесь, что GPU, выделенные для обучения и генерации, различны, чтобы избежать конфликтов ресурсов. Например, если вы планируете использовать 4 GPU для обучения и ещё 4 для генерации в vLLM, вы можете указать распределение GPU с помощью
CUDA_VISIBLE_DEVICES.Укажите GPU 0-3 для генерации в vLLM:
CUDA_VISIBLE_DEVICES=0,1,2,3 VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \ --weight-transfer-config '{"backend": "nccl"}' \ --logprobs-mode processed_logprobs \ --max-logprobs -1А GPU 4-7 для обучения:
CUDA_VISIBLE_DEVICES=4,5,6,7 torchrun --nproc_per_node 4 train.py
Сначала запустите сервер vLLM, выполнив следующую команду:
VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \
--weight-transfer-config '{"backend": "nccl"}' \
--logprobs-mode processed_logprobs \
--max-logprobs -1Затем запустите скрипт обучения и передайте use_vllm=True в аргументы обучения.
from trl import RLOOConfig
training_args = RLOOConfig(..., use_vllm=True, vllm_mode="server")Вы можете настроить конфигурацию сервера, передав дополнительные аргументы. Подробнее см. в разделе интеграция с vLLM.
При использовании vLLM убедитесь, что GPU, выделенные для обучения и генерации, различны, чтобы избежать конфликтов ресурсов. Например, если вы планируете использовать 4 GPU для обучения и ещё 4 для генерации в vLLM, вы можете указать распределение GPU с помощью
CUDA_VISIBLE_DEVICES.Укажите GPU 0-3 для генерации в vLLM:
CUDA_VISIBLE_DEVICES=0,1,2,3 VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \ --weight-transfer-config '{"backend": "nccl"}' \ --logprobs-mode processed_logprobs \ --max-logprobs -1А GPU 4-7 для обучения:
CUDA_VISIBLE_DEVICES=4,5,6,7 torchrun --nproc_per_node 4 train.py
Оптимизированные реализации внимания
TRL поддерживает различные оптимизированные реализации внимания, которые могут значительно ускорить обучение и снизить потребление памяти. Вы можете использовать как предварительно оптимизированные ядра непосредственно из Kernels Hub, так и вручную настроенный бэкенд внимания.
Вы можете использовать предварительно оптимизированные ядра внимания из Hub без ручной компиляции:
from trl import SFTConfig
training_args = SFTConfig(..., model_init_kwargs={"attn_implementation": "kernels-community/flash-attn2"})Среди других вариантов — kernels-community/vllm-flash-attn3 и kernels-community/paged-attention.
Оптимизированное внимание работает во всех тренерах TRL. Дополнительные сведения см. в разделе Kernels Hub Integration.
Ручная разработка оптимизированных бэкендов внимания — сложный и трудоёмкий процесс. Это не рекомендуется делать, если в этом нет крайней необходимости. Вместо этого рассмотрите использование Kernels из Hub, как описано в предыдущем разделе.
Если вы вручную установили оптимизированный бэкенд внимания, например Flash Attention 2, вы можете указать его в аргументах обучения:
from trl import SFTConfig
training_args = SFTConfig(..., model_init_kwargs={"attn_implementation": "flash_attention_2"})Обучение с смешанной точностью
Обучение с использованием смешанной точности bf16 или fp16 может ускорить процесс обучения и снизить потребление памяти при минимальном влиянии на качество модели.
from trl import SFTConfig
training_args = SFTConfig(..., bf16=True) # or fp16=True for older GPUsИспользуйте bf16=True для GPU Ampere (A100, RTX 30xx) и новее, а fp16=True — для более старых GPU. Обучение с смешанной точностью поддерживается во всех тренерах TRL.