Transformers documentation
Hardware
Hardware
Transformers runs on CPUs, NVIDIA GPUs, AMD GPUs, Apple Silicon, Google TPUs, AWS Trainium instances, and more through their PyTorch backends.
- For training, TrainingArguments detects the backend automatically.
- For inference, automatic detection doesn’t cover every backend, so pass the device string to choose one explicitly. This page lists the device string and package for each backend, what TrainingArguments configures it for, and how Transformers decides which device to use.
Find your hardware
Each backend has a device string that you pass to device_map or device. CPU and Apple Silicon work with the standard torch package. The other backends require a PyTorch extension package that registers the backend with PyTorch.
| Hardware | Device string | Package |
|---|---|---|
| CPU | cpu | torch |
| NVIDIA GPU | cuda | torch |
| AMD GPU | cuda | torch (ROCm build) |
| AWS Trainium | neuron | torch-neuronx |
| Google TPU | tpu | torch-tpu |
| Intel GPU | xpu | torch |
| Apple Silicon | mps | torch |
| Intel Gaudi | hpu | habana-torch-plugin |
| Huawei Ascend NPU | npu | torch-npu |
| Cambricon MLU | mlu | torch-mlu |
| Moore Threads GPU | musa | torch-musa |
If a hardware backend is missing from this table and you believe Transformers supports it, please open an issue.
How Transformers picks a device
Transformers picks a device differently for training and inference.
Training
TrainingArguments detects the backend automatically through Accelerate. In a distributed run, each process uses the device Accelerate assigns to it. In a single process, TrainingArguments picks the first available backend, and falls back to CPU if none is available. Set use_cpu=True to force training on CPU.
Inference
from_pretrained() loads a model on CPU by default.
Set device_map="auto" to place it on the available accelerators. Automatic placement relies on Accelerate to detect CUDA, Intel XPU, Intel Gaudi (HPU), and Apple Silicon (MPS), among others, but it doesn’t cover every backend. For other backends, such as AWS Trainium, pass the device string to device_map explicitly.
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B", device_map="neuron")