NVIDIA GPU

GPU NVIDIA menjalankan PyTorch melalui CUDA, dengan string perangkat cuda. Sebagian besar kernel attention dan optimasi serving di Transformers mendukung CUDA, dan beberapa di antaranya, seperti FlashAttention-3 dan FlashAttention-4, hanya berjalan di CUDA.

Instalasi

Instal driver CUDA untuk GPU Anda, lalu instal PyTorch dan Transformers. Lihat panduan instalasi untuk perangkat NVIDIA Spark (ARM64), yang memerlukan PyTorch dari indeks paket NVIDIA.

pip install "transformers[torch]"

Periksa bahwa PyTorch dapat melihat GPU.

import torch

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

Ia mencetak True dan jumlah GPU, misalnya True 4.

Pilih dan urutkan GPU

Gunakan variabel lingkungan CUDA_VISIBLE_DEVICES untuk memilih GPU dan mengatur urutannya. Tetapkan variabel ini pada baris perintah yang sama dengan eksekusi, karena variabel yang diekspor namun terlupakan dapat secara diam-diam mengarahkan eksekusi berikutnya ke GPU yang salah.

Sebagai contoh, untuk memilih GPU 0 dan 2 dari empat GPU:

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

PyTorch hanya melihat GPU 0 dan 2, yang dipetakan ke cuda:0 dan cuda:1. Untuk membalik urutannya (gunakan GPU 2 sebagai cuda:0 dan GPU 0 sebagai cuda:1):

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

Untuk menjalankan tanpa GPU:

CUDA_VISIBLE_DEVICES= python trainer-program.py ...

CUDA secara default menomori GPU dari yang tercepat, yang mungkin tidak sesuai dengan nvidia-smi. Atur CUDA_DEVICE_ORDER=PCI_BUS_ID agar indeks di CUDA_VISIBLE_DEVICES sesuai dengan nvidia-smi.

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

Pilih backend perhatian

Anda mengatur backend attention dengan attn_implementation saat memuat model, dan ini berlaku untuk inferensi maupun pelatihan. Model menggunakan SDPA secara default. GPU NVIDIA juga mendukung tiga versi FlashAttention. Jika paket FlashAttention tidak terpasang, Transformers akan memuat kernel prebuilt dari Hub, yang memerlukan 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",
)

Jalankan inferensi

Lewatkan device_map="auto" ke from_pretrained() untuk membagi model yang tidak muat di satu GPU ke beberapa GPU yang tersedia, atau device_map="cuda" untuk memuatnya di satu GPU. pipeline() memilih GPU secara otomatis jika Anda tidak melewatkan 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 mendukung CUDA graphs dan asynchronous batching pada GPU NVIDIA, serta secara otomatis menentukan apakah fitur tersebut akan digunakan.

Gunakan kernel Hub

Setel use_kernels=True di from_pretrained() untuk mengganti lapisan yang didukung dengan kernel yang dioptimalkan dari Hub. GPU NVIDIA memiliki cakupan paling luas, termasuk activations, RMSNorm, rotary embeddings, lapisan mixture-of-experts MegaBlocks, dan operasi state-space Mamba. Instal extra kernels untuk menggunakannya.

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

Hub kernels juga berlaku untuk pelatihan. Memanggil model.train() atau model.eval() akan mengganti kernel untuk mode tersebut. Untuk memilih kernel secara manual, berikan KernelConfig dan gunakan "cuda" sebagai kunci perangkat untuk pemetaan khusus NVIDIA. Lihat Memuat kernel untuk pemetaan lengkapnya.

Kuantisasi model

Quantization memungkinkan model yang lebih besar muat di memori GPU. Sebagian besar metode quantization di Transformers mendukung CUDA. Periksa kolom CUDA di ringkasan quantization untuk membandingkannya.

Untuk melakukan fine-tuning pada model yang tidak muat di memori GPU, muat model tersebut dalam 4-bit dengan bitsandbytes dan latih adapter LoRA dengan PEFT.

Latih dengan Trainer

TrainingArguments mendeteksi GPU NVIDIA dan mengatur perangkat ke cuda, sehingga Trainer dapat melatih pada GPU tersebut tanpa perlu pengaturan perangkat apa pun. Gunakan presisi campuran bf16 pada GPU Ampere atau yang lebih baru, dan fp16 pada GPU yang lebih lama.

from transformers import TrainingArguments

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

Ia mencetak cuda:0 ketika Trainer akan melatih pada satu GPU.

Pelatihan terdistribusi

Pelatihan multi-GPU menggunakan backend kolektif NCCL. Jalankan satu proses per GPU dengan torchrun atau accelerate launch.

torchrun --nproc_per_node=4 train.py

Lihat DDP, FSDP2, DeepSpeed ZeRO, dan Tensor parallelism untuk memilih strategi paralelisme.

Langkah selanjutnya

• Baca catatan semantik CUDA PyTorch untuk memahami cara PyTorch mengelola perangkat, stream, dan memori.
• Baca penggunaan memori GPU untuk memahami ke mana memori GPU digunakan selama pelatihan.
• Baca Backend perhatian untuk mempelajari bagaimana attn_implementation memilih fungsi perhatian.