Обучение на GPU AMD

Тренеры TRL работают на GPU AMD Instinct без изменения кода. Сборка PyTorch для ROCm отображает их как устройства cuda, поэтому команды из остальной части документации работают без изменений.

Установка

Установите сборку PyTorch для ROCm, затем TRL в том же окружении:

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

Ядра внимания

Ядра FlashAttention, используемые в примерах этой документации (kernels-community/flash-attn2, flash-attn3 и flash-attn4), не имеют сборки для ROCm. На GPU AMD используйте kernels-community/aiter-flash-attn, загружаемое из Hub библиотекой kernels (pip install trl[kernels]):

from trl import SFTConfig

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

Ядро не поддерживает мягкое ограничение (soft-capping) для логитов внимания, а Transformers пропускает attn_logit_softcapping без предупреждения, поэтому модели, использующие эту функцию, такие как Gemma 2, обучаются с неограниченным вниманием.

Генерация с vLLM

GRPO и RLOO с use_vllm=True требуют сборки vLLM для ROCm, см. руководство по установке vLLM.