GPU NVIDIA

Les GPU NVIDIA exécutent PyTorch via CUDA, avec la chaîne de périphérique cuda. La plupart des noyaux d'attention et des optimisations de service dans Transformers prennent en charge CUDA, et certains d'entre eux, comme FlashAttention-3 et FlashAttention-4, ne s'exécutent que sur CUDA.

Installation

Installez les pilotes CUDA pour votre GPU, puis installez PyTorch et Transformers. Consultez le guide d’installation pour les appareils NVIDIA Spark (ARM64), qui nécessitent PyTorch depuis l’index de paquets de NVIDIA.

pip install "transformers[torch]"

Vérifiez que PyTorch peut détecter les GPU.

import torch

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

Il affiche True et le nombre de GPU, par exemple True 4.

Sélectionner et ordonner les GPU

Utilisez la variable d’environnement CUDA_VISIBLE_DEVICES pour sélectionner les GPU et définir leur ordre. Définissez-la sur la même ligne de commande que l’exécution, car une variable exportée que vous oubliez peut envoyer silencieusement une exécution ultérieure vers les mauvais GPU.

Par exemple, pour sélectionner les GPU 0 et 2 sur quatre :

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

PyTorch ne voit que les GPU 0 et 2, qui sont mappés sur cuda:0 et cuda:1. Pour inverser l'ordre (utiliser le GPU 2 comme cuda:0 et le GPU 0 comme cuda:1) :

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

Pour exécuter sans aucun GPU :

CUDA_VISIBLE_DEVICES= python trainer-program.py ...

Par défaut, CUDA numérote les GPU du plus rapide au plus lent, ce qui peut ne pas correspondre à nvidia-smi. Définissez CUDA_DEVICE_ORDER=PCI_BUS_ID afin que les indices de CUDA_VISIBLE_DEVICES correspondent à nvidia-smi.

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

Choisir un backend d'attention

Vous définissez le backend d’attention avec attn_implementation lors du chargement d’un modèle, et il s’applique à la fois à l’inférence et à l’entraînement. Les modèles utilisent SDPA par défaut. Les GPU NVIDIA prennent également en charge trois versions de FlashAttention. Lorsque le package FlashAttention n’est pas installé, Transformers charge un kernel précompilé depuis le Hub, ce qui nécessite l’extra kernels (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",
)

Exécuter l'inférence

Passez device_map="auto" à from_pretrained() pour répartir un modèle qui ne tient pas sur une seule GPU sur les GPU disponibles, ou device_map="cuda" pour le charger sur une seule GPU. pipeline() sélectionne automatiquement une GPU lorsque vous ne passez pas 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 prend en charge les graphes CUDA et le groupement asynchrone sur les GPU NVIDIA, et décide automatiquement s'il doit les utiliser.

Utiliser les noyaux du Hub

Définissez use_kernels=True dans from_pretrained() pour remplacer les couches prises en charge par des noyaux optimisés issus du Hub. Les GPU NVIDIA offrent la couverture la plus large, incluant les activations, RMSNorm, les embeddings rotatifs, les couches mixture-of-experts MegaBlocks et les opérations d'espace d'états Mamba. Installez l'extra kernels pour les utiliser.

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

Les noyaux du Hub s'appliquent également à l'entraînement. L'appel de model.train() ou de model.eval() bascule sur les noyaux correspondant à ce mode. Pour choisir vous-même un noyau, passez un KernelConfig et utilisez "cuda" comme clé d'appareil pour les mappages spécifiques à NVIDIA. Consultez Chargement des noyaux pour le mappage complet.

Quantifier les modèles

La quantification permet à un modèle plus volumineux de tenir en mémoire GPU. La plupart des méthodes de quantification de Transformers prennent en charge CUDA. Consultez la colonne CUDA dans l'aperçu de la quantification pour les comparer.

Pour affiner un modèle qui ne tient pas dans la mémoire GPU, chargez-le en 4 bits avec bitsandbytes et entraînez des adaptateurs LoRA avec PEFT.

Entraîner avec Trainer

TrainingArguments détecte les GPU NVIDIA et définit l’appareil sur cuda, de sorte que Trainer entraîne sur ces GPU sans aucune configuration d’appareil. Utilisez la précision mixte bf16 sur les GPU Ampere ou plus récents, et fp16 sur les anciens.

from transformers import TrainingArguments

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

Il affiche cuda:0 lorsque Trainer entraîne sur un seul GPU.

Entraînement distribué

L’entraînement multi-GPU utilise le backend collectif NCCL. Lancez un processus par GPU avec torchrun ou accelerate launch.

torchrun --nproc_per_node=4 train.py

Consultez DDP, FSDP2, DeepSpeed ZeRO et Tensor parallelism pour choisir une stratégie de parallélisme.

Étapes suivantes

• Consultez les notes sur les semantiques CUDA de PyTorch pour comprendre comment PyTorch gère les appareils, les flux et la mémoire.
• Consultez utilisation de la mémoire GPU pour comprendre où va la mémoire GPU pendant l’entraînement.
• Lisez Backends d'attention pour comprendre comment attn_implementation sélectionne une fonction d'attention.