GPU da NVIDIA

As GPUs da NVIDIA executam o PyTorch por meio do CUDA, com a string de dispositivo cuda. A maioria dos kernels de atenção e das otimizações de serving no Transformers suporta CUDA, e alguns deles, como FlashAttention-3 e FlashAttention-4, só rodam em CUDA.

Instalação

Instale os drivers CUDA para sua GPU e, em seguida, instale o PyTorch e o Transformers. Consulte o guia de instalação para dispositivos NVIDIA Spark (ARM64), que exigem o PyTorch a partir do índice de pacotes da NVIDIA.

pip install "transformers[torch]"

Verifique se o PyTorch consegue visualizar as GPUs.

import torch

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

Ele imprime True e o número de GPUs, por exemplo True 4.

Selecionar e ordenar GPUs

Use a variável de ambiente CUDA_VISIBLE_DEVICES para selecionar as GPUs e definir a ordem delas. Defina-a na mesma linha de comando da execução, pois uma variável exportada que você esquece pode, silenciosamente, direcionar uma execução posterior para as GPUs erradas.

Por exemplo, para selecionar as GPUs 0 e 2 de um total de quatro:

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

O PyTorch enxerga apenas as GPUs 0 e 2, que são mapeadas para cuda:0 e cuda:1. Para inverter a ordem (usar a GPU 2 como cuda:0 e a GPU 0 como cuda:1):

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

Para executar sem GPUs:

CUDA_VISIBLE_DEVICES= python trainer-program.py ...

Por padrão, o CUDA numera as GPUs da mais rápida para a mais lenta, o que pode não corresponder ao nvidia-smi. Defina CUDA_DEVICE_ORDER=PCI_BUS_ID para que os índices em CUDA_VISIBLE_DEVICES correspondam ao nvidia-smi.

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

Escolha um backend de atenção

Você define o backend de atenção com attn_implementation ao carregar um modelo, e ele se aplica tanto à inferência quanto ao treinamento. Os modelos usam SDPA por padrão. As GPUs da NVIDIA também suportam três versões do FlashAttention. Quando o pacote FlashAttention não está instalado, o Transformers carrega um kernel pré-compilado do Hub, o que requer o 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",
)

Executar inferência

Passe device_map="auto" para from_pretrained() para dividir um modelo que não cabe em uma única GPU entre as GPUs disponíveis, ou device_map="cuda" para carregá-lo em uma única GPU. pipeline() seleciona uma GPU automaticamente quando você não passa 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 suporta CUDA graphs e batching assíncrono em GPUs NVIDIA e decide automaticamente se deve usá-los.

Use kernels do Hub

Defina use_kernels=True em from_pretrained() para substituir camadas compatíveis por kernels otimizados do Hub. As GPUs da NVIDIA têm a maior cobertura, incluindo ativações, RMSNorm, embeddings rotativos, camadas mixture-of-experts do MegaBlocks e operações de espaço de estado do Mamba. Instale o extra kernels para usá-los.

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

Os kernels do Hub também se aplicam ao treinamento. Chamar model.train() ou model.eval() alterna para os kernels correspondentes a cada modo. Para escolher um kernel manualmente, passe um KernelConfig e use "cuda" como chave de dispositivo para mapeamentos específicos da NVIDIA. Consulte Carregando kernels para o mapeamento completo.

Quantizar modelos

A quantização permite que um modelo maior caiba na memória da GPU. A maioria dos métodos de quantização no Transformers suporta CUDA. Consulte a coluna CUDA na visão geral da quantização para compará-los.

Para ajustar um modelo que não cabe na memória da GPU, carregue-o em 4 bits com bitsandbytes e treine adaptadores LoRA com PEFT.

Treinar com o Trainer

TrainingArguments detecta GPUs da NVIDIA e define o dispositivo como cuda, de modo que o Trainer treina neles sem qualquer configuração de dispositivo. Use precisão mista bf16 em GPUs Ampere ou mais novas, e fp16 em GPUs mais antigas.

from transformers import TrainingArguments

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

Ele imprime cuda:0 quando o Trainer vai treinar em uma única GPU.

Treinamento distribuído

O treinamento multi-GPU usa o backend coletivo NCCL. Inicie um processo por GPU com torchrun ou accelerate launch.

torchrun --nproc_per_node=4 train.py

Veja DDP, FSDP2, DeepSpeed ZeRO e Tensor parallelism para escolher uma estratégia de paralelismo.

Próximas etapas

• Leia as notas sobre semânticas de CUDA do PyTorch para entender como o PyTorch gerencia dispositivos, streams e memória.
• Leia uso de memória de GPU para entender para onde vai a memória de GPU durante o treinamento.
• Leia Backends de atenção para saber como attn_implementation seleciona uma função de atenção.