GPU de NVIDIA

Las GPU de NVIDIA ejecutan PyTorch a través de CUDA, con la cadena de dispositivo cuda. La mayoría de los núcleos de atención y las optimizaciones de servicio en Transformers admiten CUDA, y algunos de ellos, como FlashAttention-3 y FlashAttention-4, solo se ejecutan en CUDA.

Instalación

Instale los controladores de CUDA para su GPU y, a continuación, instale PyTorch y Transformers. Consulte la guía de instalación para dispositivos NVIDIA Spark (ARM64), que requieren PyTorch desde el índice de paquetes de NVIDIA.

pip install "transformers[torch]"

Comprueba que PyTorch puede ver las GPU.

import torch

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

Imprime True y el número de GPUs, por ejemplo True 4.

Seleccionar y ordenar GPUs

Use la variable de entorno CUDA_VISIBLE_DEVICES para seleccionar las GPU y establecer su orden. Defínela en la misma línea de comandos que la ejecución, ya que una variable exportada que se olvida puede enviar silenciosamente una ejecución posterior a las GPU incorrectas.

Por ejemplo, para seleccionar las GPUs 0 y 2 de un total de cuatro:

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

PyTorch solo ve las GPU 0 y 2, que se mapean a cuda:0 y cuda:1. Para invertir el orden (usar la GPU 2 como cuda:0 y la GPU 0 como cuda:1):

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

Para ejecutar sin GPUs:

CUDA_VISIBLE_DEVICES= python trainer-program.py ...

Por defecto, CUDA numera las GPU de la más rápida a la más lenta, lo cual puede no coincidir con nvidia-smi. Establezca CUDA_DEVICE_ORDER=PCI_BUS_ID para que los índices de CUDA_VISIBLE_DEVICES coincidan con nvidia-smi.

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

Elegir un backend de atención

Configura el backend de atención con attn_implementation al cargar un modelo, y se aplica tanto a la inferencia como al entrenamiento. Los modelos usan SDPA de forma predeterminada. Las GPU de NVIDIA también admiten tres versiones de FlashAttention. Cuando el paquete FlashAttention no está instalado, Transformers carga un kernel precompilado desde el Hub, lo cual requiere el 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",
)

Ejecutar inferencia

Pasa device_map="auto" a from_pretrained() para dividir un modelo que no cabe en una sola GPU entre las GPUs disponibles, o device_map="cuda" para cargarlo en una sola GPU. pipeline() elige una GPU automáticamente si no pasas 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 admite gráficos de CUDA y lotes asíncronos en GPUs de NVIDIA, y decide automáticamente si usarlos.

Usar kernels de Hub

Establezca use_kernels=True en from_pretrained() para reemplazar las capas compatibles con kernels optimizados del Hub. Las GPU de NVIDIA tienen la cobertura más amplia, que incluye activaciones, RMSNorm, embeddings rotatorios, capas mixture-of-experts de MegaBlocks y operaciones de espacio de estado de Mamba. Instale el extra kernels para utilizarlos.

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

Los kernels de Hub también se aplican al entrenamiento. Al llamar a model.train() o model.eval(), se intercambian los kernels para ese modo. Para elegir un kernel manualmente, pase un KernelConfig y use "cuda" como clave de dispositivo para las mapeos específicos de NVIDIA. Consulte Carga de kernels para el mapeo completo.

Cuantizar modelos

La cuantización permite que un modelo más grande quepa en la memoria de la GPU. La mayoría de los métodos de cuantización en Transformers admiten CUDA. Consulta la columna CUDA en el resumen de cuantización para compararlos.

Para ajustar un modelo que no cabe en la memoria de la GPU, cárgalo en 4 bits con bitsandbytes y entrena adaptadores LoRA con PEFT.

Entrenar con Trainer

TrainingArguments detecta las GPU de NVIDIA y establece el dispositivo en cuda, por lo que Trainer entrena en ellas sin necesidad de configurar el dispositivo. Utilice precisión mixta bf16 en GPU Ampere o más recientes, y fp16 en las más antiguas.

from transformers import TrainingArguments

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

Imprime cuda:0 cuando Trainer entrenará en una sola GPU.

Entrenamiento distribuido

El entrenamiento multi-GPU utiliza el backend colectivo NCCL. Inicie un proceso por GPU con torchrun o accelerate launch.

torchrun --nproc_per_node=4 train.py

Consulte DDP, FSDP2, DeepSpeed ZeRO y Tensor parallelism para elegir una estrategia de paralelismo.

Próximos pasos

• Lea las notas sobre semánticas de CUDA de PyTorch para saber cómo PyTorch gestiona los dispositivos, las streams y la memoria.
• Lee uso de memoria de GPU para entender en qué se consume la memoria de GPU durante el entrenamiento.
• Lee Backends de atención para aprender cómo attn_implementation selecciona una función de atención.