# 在 AMD GPU 上进行训练

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

## 安装

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

```bash
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`](https://huggingface.co/kernels-community/aiter-flash-attn)，它由 [`kernels`](https://github.com/huggingface/kernels) 库（`pip install trl[kernels]`）从 Hub 下载：

```python
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 安装指南](https://docs.vllm.ai/en/latest/getting_started/installation/gpu.html)。