GPU NVIDIA

GPU NVIDIA chạy PyTorch thông qua CUDA, với chuỗi thiết bị là cuda. Hầu hết các kernel attention và tối ưu hóa phục vụ trong Transformers đều hỗ trợ CUDA, và một số trong đó, chẳng hạn như FlashAttention-3 và FlashAttention-4, chỉ chạy trên CUDA.

Cài đặt

Cài đặt trình điều khiển CUDA cho GPU của bạn, sau đó cài đặt PyTorch và Transformers. Xem hướng dẫn cài đặt cho các thiết bị NVIDIA Spark (ARM64), những thiết bị này cần PyTorch từ chỉ mục gói của NVIDIA.

pip install "transformers[torch]"

Kiểm tra xem PyTorch có nhìn thấy các GPU hay không.

import torch

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

Nó in ra True và số lượng GPU, ví dụ True 4.

Chọn và sắp xếp GPU

Sử dụng biến môi trường CUDA_VISIBLE_DEVICES để chọn GPU và đặt thứ tự của chúng. Hãy đặt nó trên cùng một dòng lệnh với lệnh chạy, vì một biến đã được xuất mà bạn quên mất có thể âm thầm chuyển một lần chạy sau đó sang các GPU sai.

Ví dụ, để chọn GPU 0 và 2 trong số bốn GPU:

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

PyTorch chỉ nhìn thấy GPU 0 và 2, được ánh xạ lần lượt đến cuda:0 và cuda:1. Để đảo ngược thứ tự (dùng GPU 2 làm cuda:0 và GPU 0 làm cuda:1):

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

Để chạy mà không cần bất kỳ GPU nào:

CUDA_VISIBLE_DEVICES= python trainer-program.py ...

Mặc định, CUDA đánh số GPU theo thứ tự nhanh nhất trước, điều này có thể không khớp với nvidia-smi. Hãy đặt CUDA_DEVICE_ORDER=PCI_BUS_ID để các chỉ số trong CUDA_VISIBLE_DEVICES khớp với nvidia-smi.

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

Chọn một backend chú ý

Bạn đặt backend chú ý bằng attn_implementation khi tải một model, và nó áp dụng cho cả suy luận lẫn huấn luyện. Các model sử dụng SDPA theo mặc định. GPU NVIDIA cũng hỗ trợ ba phiên bản FlashAttention. Khi gói FlashAttention chưa được cài đặt, Transformers sẽ tải một kernel đã được xây dựng sẵn từ Hub, điều này yêu cầu 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",
)

Chạy suy luận

Truyền device_map="auto" vào from_pretrained() để chia một model không vừa trên một GPU duy nhất ra các GPU khả dụng, hoặc device_map="cuda" để nạp nó lên một GPU duy nhất. pipeline() sẽ tự động chọn một GPU khi bạn không truyền 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 hỗ trợ CUDA graphs và asynchronous batching trên GPU NVIDIA, và tự động quyết định việc sử dụng chúng.

Sử dụng kernel của Hub

Đặt use_kernels=True trong from_pretrained() để thay thế các lớp được hỗ trợ bằng các kernel tối ưu từ Hub. GPU NVIDIA có phạm vi bao phủ rộng nhất, bao gồm activations, RMSNorm, rotary embeddings, các lớp mixture-of-experts MegaBlocks và các phép toán state-space Mamba. Cài đặt extra kernels để sử dụng chúng.

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

Hub kernels cũng áp dụng cho quá trình huấn luyện. Việc gọi model.train() hoặc model.eval() sẽ thay thế các kernel tương ứng với chế độ đó. Để tự chọn một kernel, hãy truyền một KernelConfig và sử dụng "cuda" làm khóa thiết bị cho các ánh xạ riêng của NVIDIA. Xem Tải kernel để biết ánh xạ đầy đủ.

Lượng tử hóa mô hình

Quantization cho phép một mô hình lớn hơn vừa với bộ nhớ GPU. Hầu hết các phương pháp quantization trong Transformers đều hỗ trợ CUDA. Hãy kiểm tra cột CUDA trong tổng quan về quantization để so sánh chúng.

Để tinh chỉnh một mô hình không vừa trong bộ nhớ GPU, hãy tải nó ở 4-bit bằng bitsandbytes và huấn luyện các adapter LoRA bằng PEFT.

Huấn luyện với Trainer

TrainingArguments phát hiện GPU NVIDIA và đặt thiết bị thành cuda, do đó Trainer huấn luyện trên chúng mà không cần bất kỳ cài đặt thiết bị nào. Hãy dùng độ chính xác hỗn hợp bf16 trên GPU Ampere hoặc mới hơn, và fp16 trên các GPU cũ hơn.

from transformers import TrainingArguments

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

Nó in ra cuda:0 khi Trainer sẽ huấn luyện trên một GPU đơn.

Huấn luyện phân tán

Huấn luyện đa GPU sử dụng backend tập thể NCCL. Khởi chạy một tiến trình cho mỗi GPU bằng torchrun hoặc accelerate launch.

torchrun --nproc_per_node=4 train.py

Xem DDP, FSDP2, DeepSpeed ZeRO và Tensor parallelism để chọn chiến lược song song.

Các bước tiếp theo

• Đọc ghi chú về ngữ nghĩa CUDA của PyTorch để tìm hiểu cách PyTorch quản lý thiết bị, luồng và bộ nhớ.
• Đọc sử dụng bộ nhớ GPU để hiểu bộ nhớ GPU được sử dụng như thế nào trong quá trình huấn luyện.
• Đọc Các backend Attention để tìm hiểu cách attn_implementation chọn một hàm attention.