Entraînement sur des GPU AMD

Les formateurs TRL fonctionnent sur les GPU AMD Instinct sans modification de code. La version ROCm de PyTorch les expose comme des appareils cuda, de sorte que les commandes du reste de la documentation fonctionnent telles quelles.

Installation

Installez la version ROCm de PyTorch, puis TRL dans le même environnement :

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

Noyaux d'attention

Les noyaux FlashAttention utilisés dans les exemples de cette documentation (kernels-community/flash-attn2, flash-attn3 et flash-attn4) ne disposent pas de build ROCm. Sur les GPU AMD, utilisez kernels-community/aiter-flash-attn, téléchargé depuis le Hub par la bibliothèque kernels (pip install trl[kernels]) :

from trl import SFTConfig

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

Le noyau ne dispose pas de plafonnement doux des logits d'attention, et Transformers ignore attn_logit_softcapping sans avertissement, si bien que les modèles qui l'utilisent, comme Gemma 2, s'entraînent avec une attention non plafonnée.

Génération avec vLLM

GRPO et RLOO avec use_vllm=True nécessitent la version ROCm de vLLM, voir le guide d’installation de vLLM.