TRL documentation
Acelerando o Treinamento
Acelerando o Treinamento
Este guia aborda diversos métodos para acelerar o treinamento no TRL. Cada técnica inclui exemplos mínimos com links para uma documentação mais abrangente.
vLLM para geração rápida em métodos online
Métodos online, como GRPO ou Online DPO, exigem que o modelo gere complementações, o que é frequentemente um processo lento e pode impactar significativamente o tempo de treinamento. Para acelerar a geração, você pode usar vLLM, uma biblioteca que permite geração rápida por meio, entre outras coisas, do PagedAttention. Os treinadores online do TRL suportam vLLM, melhorando consideravelmente a velocidade de treinamento. Para mais detalhes, consulte Integração com vLLM.
Para usar o vLLM, primeiro instale-o usando:
pip install trl[vllm]
Primeiro, inicie um servidor vLLM executando:
VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \
--weight-transfer-config '{"backend": "nccl"}' \
--logprobs-mode processed_logprobs \
--max-logprobs -1Em seguida, execute o script de treinamento e passe use_vllm=True nos argumentos de treinamento.
from trl.experimental.online_dpo import OnlineDPOConfig
training_args = OnlineDPOConfig(..., use_vllm=True, vllm_mode="server")Primeiro, inicie um servidor vLLM executando:
VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \
--weight-transfer-config '{"backend": "nccl"}' \
--logprobs-mode processed_logprobs \
--max-logprobs -1Em seguida, execute o script de treinamento e passe use_vllm=True nos argumentos de treinamento.
from trl import GRPOConfig
training_args = GRPOConfig(..., use_vllm=True, vllm_mode="server")Você pode personalizar a configuração do servidor passando argumentos adicionais. Para mais informações, consulte integração do vLLM.
Ao usar o vLLM, certifique-se de que as GPUs atribuídas para treinamento e geração sejam separadas para evitar conflitos de recursos. Por exemplo, se você planeja usar 4 GPUs para treinamento e outras 4 para geração com o vLLM, pode especificar a alocação de GPUs usando
CUDA_VISIBLE_DEVICES.Defina as GPUs 0-3 para a geração do 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 -1E as GPUs 4-7 para o treinamento:
CUDA_VISIBLE_DEVICES=4,5,6,7 accelerate launch train.py
Primeiro, inicie um servidor vLLM executando:
VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \
--weight-transfer-config '{"backend": "nccl"}' \
--logprobs-mode processed_logprobs \
--max-logprobs -1Em seguida, execute o script de treinamento e passe use_vllm=True nos argumentos de treinamento.
from trl import RLOOConfig
training_args = RLOOConfig(..., use_vllm=True, vllm_mode="server")Você pode personalizar a configuração do servidor passando argumentos adicionais. Para mais informações, consulte integração do vLLM.
Ao usar o vLLM, certifique-se de que as GPUs atribuídas para treinamento e geração sejam separadas para evitar conflitos de recursos. Por exemplo, se você planeja usar 4 GPUs para treinamento e outras 4 para geração com o vLLM, pode especificar a alocação de GPUs usando
CUDA_VISIBLE_DEVICES.Defina as GPUs 0-3 para a geração do 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 -1E as GPUs 4-7 para o treinamento:
CUDA_VISIBLE_DEVICES=4,5,6,7 accelerate launch train.py
Implementações de atenção otimizadas
O TRL suporta várias implementações de atenção otimizadas que podem acelerar significativamente o treinamento, ao mesmo tempo em que reduzem o uso de memória. Você pode usar kernels pré-otimizados diretamente do Kernels Hub ou um backend de atenção construído manualmente.
Você pode usar kernels de atenção pré-otimizados do Hub sem compilação manual:
from trl import SFTConfig
training_args = SFTConfig(..., model_init_kwargs={"attn_implementation": "kernels-community/flash-attn2"})Outras opções incluem kernels-community/vllm-flash-attn3 e kernels-community/paged-attention.
A atenção otimizada funciona em todos os treinadores do TRL. Para mais detalhes, consulte Integração com o Kernels Hub.
Construir manualmente backends de atenção otimizados é complexo e demorado. Não é recomendado a menos que seja absolutamente necessário. Considere usar Kernels do Hub, conforme descrito na seção anterior.
Se você instalou manualmente um backend de atenção otimizado, como Flash Attention 2, você pode especificá-lo nos argumentos de treinamento:
from trl import SFTConfig
training_args = SFTConfig(..., model_init_kwargs={"attn_implementation": "flash_attention_2"})Treinamento com precisão mista
O treinamento com precisão mista usando bf16 ou fp16 pode acelerar o treinamento e reduzir o uso de memória com impacto mínimo na qualidade do modelo.
from trl import SFTConfig
training_args = SFTConfig(..., bf16=True) # or fp16=True for older GPUsUse bf16=True para GPUs Ampere (A100, RTX 30xx) ou mais novas, e fp16=True para GPUs mais antigas. O treinamento em precisão mista é suportado em todos os treinadores do TRL.