NVIDIA GPU

NVIDIA GPUs работают с PyTorch через CUDA, при этом строка устройства — cuda. Большинство ядер внимания и оптимизаций для инференса в Transformers поддерживают CUDA, а некоторые из них, такие как FlashAttention-3 и FlashAttention-4, работают только на CUDA.

Установка

Установите драйверы CUDA для вашего GPU, затем установите PyTorch и Transformers. См. руководство по установке для устройств NVIDIA Spark (ARM64), для которых требуется PyTorch из репозитория пакетов NVIDIA.

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. GPU от NVIDIA также поддерживают три версии FlashAttention. Если пакет FlashAttention не установлен, Transformers загружает готовое ядро из Hub, для чего требуется дополнительный пакет 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",
)

Запуск вывода

Передайте device_map="auto" в from_pretrained(), чтобы распределить модель, которая не помещается на одной GPU, по доступным GPU, или device_map="cuda", чтобы загрузить её на одну GPU. pipeline() автоматически выбирает GPU, если вы не передаёте 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 поддерживает CUDA-графы и асинхронную пакетную обработку на GPU NVIDIA и автоматически определяет, следует ли их использовать.

Использование ядер Hub

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 применяются и к обучению. Вызов model.train() или model.eval() подменяет ядра на соответствующие данному режиму. Чтобы выбрать ядро самостоятельно, передайте KernelConfig и используйте "cuda" в качестве ключа устройства для специфичных для NVIDIA отображений. Полное отображение описано в разделе Загрузка ядер.

Квантование моделей

Квантование позволяет разместить более крупную модель в памяти GPU. Большинство методов квантования в Transformers поддерживают CUDA. Сравните их, обратив внимание на столбец CUDA в обзоре квантования.

Чтобы дообучить модель, которая не помещается в память GPU, загрузите её в 4-битном формате с помощью bitsandbytes и обучите адаптеры LoRA с помощью PEFT.

Обучение с помощью Trainer

TrainingArguments определяет GPU NVIDIA и устанавливает устройство в cuda, поэтому Trainer обучает на них без каких-либо настроек устройства. Используйте смешанную точность bf16 на GPU Ampere или новее, а fp16 — на более старых.

from transformers import TrainingArguments

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

Выводит cuda:0, если Trainer обучает на одном GPU.

Распределённое обучение

Обучение на нескольких GPU использует коллективный бэкенд NCCL. Запускайте по одному процессу на каждый GPU с помощью torchrun или accelerate launch.

torchrun --nproc_per_node=4 train.py

См. DDP, FSDP2, DeepSpeed ZeRO и Tensor parallelism, чтобы выбрать стратегию параллелизма.

Дальнейшие шаги

• Ознакомьтесь с примечаниями о семантике CUDA в PyTorch, чтобы узнать, как PyTorch управляет устройствами, потоками и памятью.
• Ознакомьтесь с разделом Использование памяти GPU, чтобы понять, куда уходит память GPU во время обучения.
• Ознакомьтесь с разделом Бэкенды внимания, чтобы узнать, как attn_implementation выбирает функцию внимания.