AMD GPU

AMD Instinct（MI200、MI300、MI350 系列）和 Radeon GPU 通过 AMD 的 GPU 软件栈 ROCm 运行 PyTorch。PyTorch 的 ROCm 版本使用 AMD 的类 CUDA 编程接口 HIP，通过 torch.cuda API 暴露 AMD GPU。设备字符串为 cuda，因此为 NVIDIA GPU 编写的代码可以无需修改直接运行。Transformers 在运行时检查 torch.version.hip 以区分两者，并将 attention 和 Hub 内核路由到 AMD 专用构建。

安装

从 huggingface/transformers-pytorch-amd-gpu Docker 镜像安装 ROCm。该镜像附带了 PyTorch 的 ROCm 构建版本、适用于 MI300 和 MI350 GPU（gfx942 和 gfx950）的 flash-attn 构建版本，以及从源码安装的 Transformers。

docker run -it --device=/dev/kfd --device=/dev/dri --group-add video \
    --ipc=host --shm-size 16G \
    huggingface/transformers-pytorch-amd-gpu

--device=/dev/kfd 和 --device=/dev/dri 会将 AMD GPU 暴露给容器。在容器内安装 kernels 以 使用 Hub kernels。

pip install "transformers[kernels]"

若要安装到您自己的系统上，请在 PyTorch 安装选择器 中选择 ROCm，并确保 wheel 的 ROCm 版本受您的驱动程序支持。

检查 PyTorch 是否识别到 AMD GPU。

import torch

print(torch.version.hip, torch.cuda.is_available(), torch.cuda.device_count())
print(torch.cuda.get_device_properties(0).gcnArchName)

7.2.53211 True 2
gfx942:sramecc+:xnack-

在 PyTorch 的 CUDA 构建中，torch.version.hip 的值为 None。gcnArchName 会报告 GPU 架构，例如 MI300 对应 gfx942，MI350 对应 gfx950，Ryzen AI Max (Strix Halo) 对应 gfx1151。

选择和排序 GPU

使用 HIP_VISIBLE_DEVICES 来选择 GPU 并设置其顺序。请在与运行命令相同的命令行上设置该变量。

HIP_VISIBLE_DEVICES=0,2 torchrun trainer-program.py ...

PyTorch 只能看到 GPU 0 和 2，它们分别映射到 cuda:0 和 cuda:1。ROCR_VISIBLE_DEVICES 在更底层的堆栈中过滤 GPU，通常由集群调度器设置。如果两者都设置了，HIP_VISIBLE_DEVICES 会对 ROCR_VISIBLE_DEVICES 保留的可见 GPU 进行索引。

使用 rocm-smi 或 amd-smi 列出 GPU 及其索引。

运行推理

Pipeline 和 device_map="auto" 会自动检测 AMD GPU 并将模型放置在 cuda 上。您也可以显式传入 device_map="cuda" 或 device_map="cuda:1"。

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B", dtype="bfloat16", device_map="auto"
)

inputs = tokenizer("The French Bread Law states", return_tensors="pt").to(model.device)
output = model.generate(**inputs, do_sample=False, max_new_tokens=50)
print(tokenizer.decode(output[0], skip_special_tokens=True))

The French Bread Law states that the French government must provide a certain amount of bread for the citizens of the country. The law is passed by the French parliament. The French parliament is a body of the French government. The French government is a body of the French state. The

选择注意力后端

在加载模型时，使用 attn_implementation 设置注意力后端，该设置同时适用于推理和训练。SDPA 是默认选项，适用于 PyTorch 支持的所有 AMD GPU。此外，还提供两种 FlashAttention 选项。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    dtype="bfloat16",
    device_map="auto",
    attn_implementation="flash_attention_3",
)

在 ROCm 上，flash_attention_3 会解析为 AITER 内核并记录以下日志。Transformers 目前使用的是 kernels-community/aiter-flash-attn 别名，该别名会重定向到 amd/aiter-flash-attn。此消息指向 FlashAttention-3 包，因此即使已安装 flash_attention_2，该消息仍会出现。

You do not have `flash_attn` installed, using `kernels-community/aiter-flash-attn` from the `kernels` library instead!

您也可以直接通过 attn_implementation="kernels-community/aiter-flash-attn" 传递内核仓库。有关后端的选择方式，请参阅 Attention backends。

使用 Hub 内核

在 from_pretrained() 中设置 use_kernels=True，即可用 Hub 上的优化内核替换受支持的层。Transformers 会检测 ROCm，并为每个层选择 AMD 构建版本，例如旋转嵌入使用 amd/aiter-rope，RMSNorm 使用 kernels-community/liger-kernels。编写自定义 KernelConfig 时，请使用 "rocm" 作为 AMD 特定映射的设备键。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B", dtype="bfloat16", device_map="auto", use_kernels=True
)

Hub 内核主要加速 prefill 阶段和大批量处理。在启用它们进行小批量解码之前，请先针对自身工作负载进行基准测试。完整的映射关系请参阅 加载内核。

量化模型

ROCm 上支持的方法包括 AWQ、GPT-QModel、compressed-tensors 以及 AMD 的 Quark。bitsandbytes 仅获得部分支持。MI300 及更新型号的 GPU 原生支持 FP8。完整的列表请参阅兼容性表。

使用 Trainer 进行训练

TrainingArguments 会检测 AMD GPU，并自动将设备设置为 cuda。

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./outputs",
    bf16=True,
)

分布式训练

PyTorch 的 ROCm 版本附带了 RCCL，即 AMD 对 NCCL API 的实现，因此 nccl 后端、torchrun、FSDP、DeepSpeed 以及 Transformers 分布式 API 均可无需修改直接使用。

import torch
from transformers import AutoModelForCausalLM, DistributedConfig

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-1.7B",
    dtype=torch.bfloat16,
    distributed_config=DistributedConfig(tp_size=2, fsdp_size=4),
)

torchrun --nproc_per_node=8 train.py

有关可用选项，请参阅 张量并行、专家并行 和 FSDP。

局限性

• CodeCarbonCallback 与 ROCm 不兼容。请在 TrainingArguments 中将 report_to 设置为排除 codecarbon。
• 在 兼容性表 中标记为 🔴 的 ROCm 量化方法依赖于仅限 CUDA 的内核，无法在 AMD GPU 上加载。

后续步骤

• 阅读 ROCm 文档以了解驱动程序和 PyTorch 的设置。
• 在 AMD 和 kernels-community 组织中浏览 ROCm 构建。
• 阅读 加载内核 以自定义哪些内核替换哪些层。