Transformers documentation
硬件
You are viewing main version, which requires installation from source. If you'd like
regular pip install, checkout the latest stable version (v5.22.1).
硬件
Transformers 通过其 PyTorch 后端支持在 CPU、NVIDIA GPU、AMD GPU、Apple Silicon、Google TPU、AWS Trainium 实例等多种硬件上运行。
- 在训练过程中，TrainingArguments 会自动检测后端。
- 在推理阶段，自动检测无法覆盖所有后端，因此请传入设备字符串以显式指定后端。本页列出了每个后端对应的设备字符串和软件包、TrainingArguments 为其配置的内容，以及 Transformers 如何决定使用哪个设备。
查找您的硬件
每个后端都有一个设备字符串，你需要将其传递给 device_map 或 device。CPU 和 Apple Silicon 可以使用标准的 torch 包。其他后端则需要一个 PyTorch 扩展包，用于将该后端注册到 PyTorch 中。
| 硬件 | 设备字符串 | 包 |
|---|---|---|
| CPU | cpu | torch |
| NVIDIA GPU | cuda | torch |
| AMD GPU | cuda | torch（ROCm 构建版本） |
| AWS Trainium | neuron | torch-neuronx |
| Google TPU | tpu | torch-tpu |
| Intel GPU | xpu | torch |
| Apple Silicon | mps | torch |
| Intel Gaudi | hpu | habana-torch-plugin |
| 华为 Ascend NPU | npu | torch-npu |
| 寒武纪 MLU | mlu | torch-mlu |
| Moore Threads GPU | musa | torch-musa |
如果此表中缺少某个硬件后端，且您认为 Transformers 支持该后端，请提交 issue。
Transformers 如何选择合适的设备
Transformers 在训练和推理时选择不同的设备。
训练
TrainingArguments 会通过 Accelerate 自动检测后端。在分布式运行中，每个进程使用 Accelerate 为其分配的设备。在单进程情况下，TrainingArguments 会选择第一个可用的后端，如果没有可用后端则回退到 CPU。设置 use_cpu=True 可强制在 CPU 上进行训练。
推理
from_pretrained() 默认在 CPU 上加载模型。
设置 device_map="auto" 可将模型放置在可用的加速器上。自动放置依赖于 Accelerate 来检测 CUDA、Intel XPU、Intel Gaudi (HPU) 和 Apple Silicon (MPS) 等后端，但它并不涵盖所有后端。对于其他后端，例如 AWS Trainium 和 Google TPU，请显式地将设备字符串传递给 device_map。
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B", device_map="neuron")