NVIDIA GPU

NVIDIA GPU 通过 CUDA 运行 PyTorch，设备字符串为 cuda。Transformers 中大多数注意力内核和推理优化都支持 CUDA，其中一些（如 FlashAttention-3 和 FlashAttention-4）仅支持 CUDA。

安装

请为您的 GPU 安装 CUDA 驱动程序，然后安装 PyTorch 和 Transformers。有关 NVIDIA Spark (ARM64) 设备的安装指南，请参阅安装指南，这些设备需要从 NVIDIA 的软件包索引中安装 PyTorch。

pip install "transformers[torch]"

检查 PyTorch 是否能够识别 GPU。

import torch

print(torch.cuda.is_available(), torch.cuda.device_count())

它会打印 True 以及 GPU 的数量，例如 True 4。

选择和排序 GPU

使用 CUDA_VISIBLE_DEVICES 环境变量来选择 GPU 并设置其顺序。请在与运行命令相同的命令行中设置该变量，因为忘记取消导出的变量可能会悄悄地将后续运行发送到错误的 GPU 上。

例如，要从四块 GPU 中选择第 0 号和第 2 号：

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

PyTorch 只能看到 GPU 0 和 2，它们分别映射到 cuda:0 和 cuda:1。若要反转顺序（将 GPU 2 用作 cuda:0，将 GPU 0 用作 cuda:1）：

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

要在没有任何 GPU 的情况下运行：

CUDA_VISIBLE_DEVICES= python trainer-program.py ...

默认情况下，CUDA 会按速度从快到慢对 GPU 进行编号，这可能与 nvidia-smi 中的顺序不一致。请设置 CUDA_DEVICE_ORDER=PCI_BUS_ID，使 CUDA_VISIBLE_DEVICES 中的索引与 nvidia-smi 保持一致。

CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES=0,2 torchrun trainer-program.py ...

选择注意力后端

在加载模型时，通过 attn_implementation 设置注意力后端，该设置同时适用于推理和训练。模型默认使用 SDPA。NVIDIA GPU 还支持三个版本的 FlashAttention。如果未安装 FlashAttention 包，Transformers 会从 Hub 加载预构建的 kernel，这需要 kernels 附加组件（pip install "transformers[kernels]"）。

`attn_implementation`	GPU	当未安装包时，使用 Hub 内核
`"flash_attention_2"`	Ampere 或更新	`kernels-community/flash-attn2`
`"flash_attention_3"`	Ampere 或更新	`kernels-community/vllm-flash-attn3`
`"flash_attention_4"`	Hopper 或更新	`kernels-community/flash-attn4`

from transformers import AutoModelForCausalLM

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

运行推理

将 device_map="auto" 传递给 from_pretrained()，可将无法装入单个 GPU 的模型拆分到可用的多个 GPU 上；或者使用 device_map="cuda" 将其加载到单个 GPU 上。如果未传递 device，pipeline() 会自动选择一个 GPU。

from transformers import AutoModelForCausalLM, AutoTokenizer

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

inputs = tokenizer("NVIDIA GPUs are", return_tensors="pt").to(model.device)
output = model.generate(**inputs, max_new_tokens=50, do_sample=False)
print(tokenizer.decode(output[0], skip_special_tokens=True))

连续批处理支持在 NVIDIA GPU 上使用 CUDA 图和异步批处理，并会自动决定是否启用它们。

使用 Hub 内核

在 from_pretrained() 中设置 use_kernels=True，即可将受支持的层替换为来自 Hub 的优化内核。NVIDIA GPU 的覆盖范围最广，包括激活函数、RMSNorm、旋转位置编码、MegaBlocks 混合专家层以及 Mamba 状态空间算子。请安装 kernels 附加组件以使用这些功能。

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

Hub 内核同样适用于训练。调用 model.train() 或 model.eval() 会切换为该模式对应的内核。若要自行选择内核，请传入 KernelConfig，并使用 "cuda" 作为设备键，以指定 NVIDIA 专属的映射。完整的映射关系请参阅加载内核。

量化模型

量化技术可以让更大的模型装入 GPU 显存。Transformers 中的大多数量化方法都支持 CUDA。请查看量化概览中的 CUDA 列以进行对比。

要微调无法装入 GPU 显存的模型，请使用 bitsandbytes 以 4-bit 精度加载模型，并使用 PEFT 训练 LoRA 适配器。

使用 Trainer 进行训练

TrainingArguments 会检测 NVIDIA GPU 并将设备设置为 cuda，因此 Trainer 无需任何设备配置即可在这些 GPU 上进行训练。在 Ampere 或更新的 GPU 上使用 bf16 混合精度，在较旧的 GPU 上使用 fp16。

from transformers import TrainingArguments

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

当 Trainer 将在单个 GPU 上进行训练时，它会打印 cuda:0。

分布式训练

多 GPU 训练使用 NCCL 集合通信后端。请使用 torchrun 或 accelerate launch 为每个 GPU 启动一个进程。

torchrun --nproc_per_node=4 train.py

请参阅 DDP、FSDP2、DeepSpeed ZeRO 和 Tensor parallelism，以选择合适的并行策略。

后续步骤

阅读 PyTorch 的 CUDA 语义说明，了解 PyTorch 如何管理设备、流和内存。
阅读GPU 内存使用情况，以了解训练过程中 GPU 内存的消耗去向。
阅读 Attention 后端以了解 attn_implementation 如何选择注意力函数。

Update on GitHub