Entrenamiento en GPUs de AMD

Los entrenadores de TRL se ejecutan en GPUs AMD Instinct sin cambios de código. La compilación ROCm de PyTorch los expone como dispositivos cuda, por lo que los comandos del resto de la documentación funcionan tal cual.

Instalación

Instala la compilación de ROCm de PyTorch y luego TRL en el mismo entorno:

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

Kernels de atención

Los kernels de FlashAttention utilizados en los ejemplos de esta documentación (kernels-community/flash-attn2, flash-attn3 y flash-attn4) no tienen una compilación para ROCm. En GPUs de AMD, use kernels-community/aiter-flash-attn, descargado del Hub por la biblioteca kernels (pip install trl[kernels]):

from trl import SFTConfig

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

El kernel no tiene soft-capping de logits de atención, y Transformers omite attn_logit_softcapping sin emitir una advertencia, por lo que los modelos que lo utilizan, como Gemma 2, se entrenan con atención sin límite.

Generación con vLLM

GRPO y RLOO con use_vllm=True necesitan la compilación de ROCm de vLLM, consulte la guía de instalación de vLLM.