在 AMD GPU 上进行训练

TRL 训练器可在 AMD Instinct GPU 上运行，无需修改代码。PyTorch 的 ROCm 版本将它们暴露为 cuda 设备，因此文档其余部分中的命令可原样使用。

安装

安装 PyTorch 的 ROCm 版本，然后在同一环境中安装 TRL：

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

Attention 内核

本文档示例中使用的 FlashAttention 内核（kernels-community/flash-attn2、flash-attn3 和 flash-attn4）没有 ROCm 构建版本。在 AMD GPU 上，请使用 kernels-community/aiter-flash-attn，它由 kernels 库（pip install trl[kernels]）从 Hub 下载：

from trl import SFTConfig

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

该内核没有注意力对数几率软截断功能，且 Transformers 在跳过 attn_logit_softcapping 时不会发出警告，因此使用此功能的模型（例如 Gemma 2）将以无截断的注意力进行训练。

使用 vLLM 进行生成

使用 use_vllm=True 的 GRPO 和 RLOO 需要 vLLM 的 ROCm 构建版本，请参阅 vLLM 安装指南。