AMD GPU에서 학습

TRL trainers는 코드 변경 없이 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 라이브러리가(pip install trl[kernels]) Hub에서 다운로드하는 kernels-community/aiter-flash-attn를 사용하세요:

from trl import SFTConfig

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

커널에는 attention-logit soft-capping 기능이 없으며, Transformers는 경고 없이 attn_logit_softcapping을 건너뜁니다. 따라서 Gemma 2와 같이 이를 사용하는 모델은 캡이 없는 attention으로 학습됩니다.

vLLM을 사용한 생성

use_vllm=True를 사용하는 GRPO와 RLOO에는 vLLM의 ROCm 빌드가 필요합니다. vLLM 설치 가이드를 참고하세요.