Treinamento em GPUs AMD

Os treinadores do TRL rodam em GPUs AMD Instinct sem alterações de código. A build ROCm do PyTorch os expõe como dispositivos cuda, então os comandos no restante da documentação funcionam como estão.

Instalação

Instale o build do PyTorch para ROCm e, em seguida, o TRL no mesmo ambiente:

pip install torch torchvision --index-url https://download.pytorch.org/whl/rocm7.2
pip install trl

Kernels de atenção

Os kernels FlashAttention usados nos exemplos desta documentação (kernels-community/flash-attn2, flash-attn3 e flash-attn4) não possuem build para ROCm. Em GPUs AMD, use kernels-community/aiter-flash-attn, baixado do Hub pela biblioteca kernels (pip install trl[kernels]):

from trl import SFTConfig

training_args = SFTConfig(model_init_kwargs={"attn_implementation": "kernels-community/aiter-flash-attn@v2"})

O kernel não possui soft-capping de logit de atenção, e o Transformers ignora attn_logit_softcapping sem um aviso, de modo que modelos que o utilizam, como o Gemma 2, são treinados com atenção sem limite.

Geração com vLLM

GRPO e RLOO com use_vllm=True precisam do build ROCm do vLLM, consulte o guia de instalação do vLLM.