NVIDIA GPU

NVIDIA GPUs run PyTorch through CUDA, with the device string cuda. Most attention kernels and serving optimizations in Transformers support CUDA, and some of them, like FlashAttention-3 and FlashAttention-4, only run on CUDA.

Installation

Install the CUDA drivers for your GPU, then install PyTorch and Transformers. See the installation guide for NVIDIA Spark (ARM64) devices, which need PyTorch from NVIDIA’s package index.

pip install "transformers[torch]"

Check that PyTorch can see the GPUs.

import torch

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

It prints True and the number of GPUs, for example True 4.

Select and order GPUs

Use the CUDA_VISIBLE_DEVICES environment variable to select GPUs and set their order. Set it on the same command line as the run, because an exported variable you forget about can silently send a later run to the wrong GPUs.

For example, to select GPUs 0 and 2 out of four:

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

PyTorch sees only GPUs 0 and 2, which are mapped to cuda:0 and cuda:1. To reverse the order (use GPU 2 as cuda:0 and GPU 0 as cuda:1):

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

To run without any GPUs:

CUDA_VISIBLE_DEVICES= python trainer-program.py ...

CUDA numbers GPUs fastest first by default, which may not match nvidia-smi. Set CUDA_DEVICE_ORDER=PCI_BUS_ID so the indices in CUDA_VISIBLE_DEVICES match nvidia-smi.

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

Choose an attention backend

You set the attention backend with attn_implementation when you load a model, and it applies to both inference and training. Models use SDPA by default. NVIDIA GPUs also support three FlashAttention versions. When the FlashAttention package isn’t installed, Transformers loads a prebuilt kernel from the Hub instead, which needs the kernels extra (pip install "transformers[kernels]").

from transformers import AutoModelForCausalLM

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

Run inference

Pass device_map="auto" to from_pretrained() to split a model that doesn’t fit on one GPU across the available GPUs, or device_map="cuda" to load it on a single GPU. pipeline() picks a GPU automatically when you don’t pass device.

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))

Continuous batching supports CUDA graphs and asynchronous batching on NVIDIA GPUs, and decides whether to use them automatically.

Use Hub kernels

Set use_kernels=True in from_pretrained() to replace supported layers with optimized kernels from the Hub. NVIDIA GPUs have the widest coverage, including activations, RMSNorm, rotary embeddings, MegaBlocks mixture-of-experts layers, and Mamba state-space ops. Install the kernels extra to use them.

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

Hub kernels apply to training too. Calling model.train() or model.eval() swaps in the kernels for that mode. To choose a kernel yourself, pass a KernelConfig and use "cuda" as the device key for NVIDIA-specific mappings. See Loading kernels for the full mapping.

Quantize models

Quantization lets a larger model fit in GPU memory. Most quantization methods in Transformers support CUDA. Check the CUDA column in the quantization overview to compare them.

To fine-tune a model that doesn’t fit in GPU memory, load it in 4-bit with bitsandbytes and train LoRA adapters with PEFT.

Train with Trainer

TrainingArguments detects NVIDIA GPUs and sets the device to cuda, so Trainer trains on them without any device settings. Use bf16 mixed precision on Ampere or newer GPUs, and fp16 on older ones.

from transformers import TrainingArguments

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

It prints cuda:0 when Trainer will train on a single GPU.

Distributed training

Multi-GPU training uses the NCCL collective backend. Launch one process per GPU with torchrun or accelerate launch.

torchrun --nproc_per_node=4 train.py

See DDP, FSDP2, DeepSpeed ZeRO, and Tensor parallelism to choose a parallelism strategy.

Next steps

• Read the PyTorch CUDA semantics notes for how PyTorch manages devices, streams, and memory.
• Read GPU memory usage to understand where GPU memory goes during training.
• Read Attention backends to learn how attn_implementation selects an attention function.