NVIDIA GPU

NVIDIA GPU는 cuda 디바이스 문자열을 사용하여 PyTorch를 CUDA를 통해 실행합니다. Transformers의 대부분의 attention 커널 및 서빙 최적화는 CUDA를 지원하며, FlashAttention-3와 FlashAttention-4와 같은 일부는 CUDA에서만 실행됩니다.

설치

GPU에 맞는 CUDA 드라이버를 설치한 후, PyTorch와 Transformers를 설치하세요. NVIDIA의 패키지 인덱스에서 PyTorch를 설치해야 하는 NVIDIA Spark (ARM64) 장치에 대한 설치 가이드를 참고하세요.

pip install "transformers[torch]"

PyTorch가 GPU를 인식할 수 있는지 확인합니다.

import torch

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

True와 GPU의 수를 출력합니다. 예를 들어 True 4와 같습니다.

GPU 선택 및 순서 지정

GPU를 선택하고 그 순서를 설정하려면 CUDA_VISIBLE_DEVICES 환경 변수를 사용하세요. 내보낸 변수를 잊어버리면 이후 실행이 조용히 잘못된 GPU로 전송될 수 있으므로, 실행과 동일한 명령줄에서 설정하세요.

예를 들어, 4개의 GPU 중 0번과 2번을 선택하려면:

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

PyTorch는 cuda:0과 cuda:1에 매핑된 GPU 0과 2만 인식합니다. 순서를 반대로 바꾸려면(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_VISIBLE_DEVICES의 인덱스가 nvidia-smi와 일치하도록 CUDA_DEVICE_ORDER=PCI_BUS_ID를 설정하세요.

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

어텐션 백엔드 선택

모델을 로드할 때 attn_implementation으로 attention back end를 설정하며, 이는 추론과 학습 모두에 적용됩니다. 모델은 기본적으로 SDPA를 사용합니다. NVIDIA GPU는 FlashAttention의 세 가지 버전을 지원합니다. FlashAttention 패키지가 설치되어 있지 않은 경우, Transformers는 Hub에서 사전 빌드된 커널을 로드하며, 이때 kernels extra가 필요합니다(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에 로드합니다. device를 전달하지 않으면 pipeline()이 GPU를 자동으로 선택합니다.

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는 NVIDIA GPU에서 CUDA graphs와 비동기 배치 처리를 지원하며, 이를 사용할지 여부를 자동으로 결정합니다.

Hub 커널 사용

use_kernels=True를 from_pretrained()에 설정하면 Hub의 최적화 커널로 지원되는 레이어를 대체합니다. NVIDIA GPU가 가장 광범위한 범위를 지원하며, 활성화 함수, RMSNorm, rotary embeddings, MegaBlocks mixture-of-experts 레이어, Mamba state-space ops를 포함합니다. 이를 사용하려면 kernels extra를 설치하세요.

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

Hub 커널은 학습에도 적용됩니다. model.train() 또는 model.eval()을 호출하면 해당 모드에 맞는 커널로 교체됩니다. 커널을 직접 선택하려면 KernelConfig를 전달하고, NVIDIA 전용 매핑의 디바이스 키로 "cuda"를 사용하세요. 전체 매핑은 커널 로딩을 참고하세요.

모델 양자화

Quantization을 사용하면 더 큰 모델을 GPU 메모리에 담을 수 있습니다. Transformers의 대부분의 양자화 방법은 CUDA를 지원합니다. 양자화 개요의 CUDA 열을 확인하여 각 방법을 비교해 보세요.

GPU 메모리에 맞지 않는 모델을 미세 조정하려면 bitsandbytes로 4비트로 로드하고 PEFT로 LoRA 어댑터를 학습하세요.

Trainer로 학습

TrainingArguments는 NVIDIA GPU를 감지하여 장치를 cuda로 설정하므로, Trainer는 별도의 장치 설정 없이 해당 GPU에서 학습을 수행합니다. Ampere 또는 그 이후의 GPU에서는 bf16 혼합 정밀도를, 이전 세대 GPU에서는 fp16을 사용하세요.

from transformers import TrainingArguments

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

Trainer가 단일 GPU에서 학습할 경우 cuda:0를 출력합니다.

분산 학습

멀티 GPU 학습은 NCCL 콜렉티브 백엔드를 사용합니다. torchrun 또는 accelerate launch를 사용하여 GPU당 하나의 프로세스를 시작하세요.

torchrun --nproc_per_node=4 train.py

DDP, FSDP2, DeepSpeed ZeRO, Tensor parallelism를 참고하여 병렬화 전략을 선택하세요.

다음 단계

• PyTorch가 장치, 스트림, 메모리를 관리하는 방법에 대해서는 PyTorch CUDA semantics 노트를 참고하세요.
• 트레이닝 중 GPU 메모리가 어디에 사용되는지 이해하려면 GPU memory usage를 읽어보세요.
• attn_implementation이 어텐션 함수를 선택하는 방식을 배우려면 Attention backends를 읽어보세요.