GPU에서 효율적인 모델 훈련
GPU에서 효율적인 모델 훈련
GPU로 딥러닝 모델을 훈련할 때 가장 많이 고민하는 게 "메모리와 속도의 균형"이에요. GPU의 높은 메모리 대역폭과 병렬 처리 능력을 잘 활용하면 수십억 파라미터 모델도 훈련할 수 있지만, GPU 메모리 활용(데이터 처리량/훈련 시간)과 훈련 속도 사이의 균형을 찾는 게 핵심입니다. 이 가이드에서는 Transformers와 PyTorch가 제공하는 효율적 GPU 훈련 기능을 정리해요.
출처: GPU training efficiency - Hugging Face Transformers 공식문서
기능 한눈에 보기
아래 표로 훈련 시나리오에 맞는 기능을 빠르게 찾을 수 있어요.
| Feature | Training speed | Memory usage |
|---|---|---|
| batch size | yes | yes |
| gradient accumulation | no | yes |
| gradient checkpointing | no | yes |
| mixed precision | yes | depends |
| optimizers | yes | yes |
| data preloading | yes | no |
| torch_empty_cache_steps | no | yes |
| torch.compile | yes | no |
| scaled dot production attention (SDPA) | yes | yes |
Trainer
Trainer는 TrainingArguments를 통해 설정할 수 있는 유용한 훈련 기능을 많이 지원합니다. 이 절에서는 훈련 최적화에 중요한 기능 몇 가지를 강조합니다.
배치 크기
배치 크기(batch size)는 메모리 사용과 훈련 속도에 영향을 주기 때문에 효율적 GPU 훈련에서 가장 중요한 하이퍼파라미터 중 하나예요. 배치가 클수록 GPU의 병렬 처리 능력을 활용해 훈련이 빨라집니다. 2의 거듭제곱(8, 64, 128, 256, 512 등) 배치 크기를 권장해요. 배치 크기는 GPU와 모델 데이터 타입에 따라 달라집니다.
TrainingArguments에서 per_device_train_batch_size()를 설정합니다.
from transformers import TrainingArguments
args = TrainingArguments(
per_device_train_batch_size=256,
per_device_eval_batch_size=256,
)
NVIDIA Performance 가이드에서 입력 특성·출력 뉴런 수·배치 크기가 성능에 미치는 영향을 배울 수 있어요. 이들은 GPU가 수행하는 일반 행렬 곱(GEMM)과 연관되어 있고, 큰 파라미터일수록 병렬화·효율에 좋습니다.
Tensor Core Requirements 절도 데이터 타입과 GPU에 따라 텐서 곱셈 속도를 최대화하는 배치 크기를 고르는 데 유용해요. 예를 들어 fp16에서는 8의 배수, A100 GPU라면 64의 배수를 권장합니다.
마지막으로 작은 파라미터에 대한 Dimension Quantization Effects를 고려하세요. 행렬 차원이 GPU 스레드 블록 타일 크기로 나누어떨어지지 않으면 타일 양자화가 발생해 GPU 자원을 덜 활용합니다. 타일 크기로 나누어떨어지는 배치 크기 배수를 고르면 훈련을 크게 빠르게 할 수 있어요.
Gradient accumulation
그래디언트 누적(gradient accumulation)은 파라미터를 갱신하기 전에 여러 미니 배치에 걸쳐 그래디언트를 누적해 메모리 제약을 극복합니다. 그래디언트를 덜 저장하므로 메모리가 줄고, 보통 단일 배치에서 파라미터가 갱신되는 점을 활용해 더 큰 effective batch size 로 훈련할 수 있어요. 다만 누적에 의한 추가 forward·backward pass 때문에 훈련이 느려질 수 있습니다.
TrainingArguments에서 per_device_train_batch_size()와 gradient_accumulation_steps를 설정합니다.
from transformers import TrainingArguments
# effective batch size of 64
args = TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=16,
)
너무 많은 누적 스텝은 훈련을 크게 느리게 만들 수 있으니 피하세요. 아래 예시에서 GPU에 맞는 최대 배치 크기가 4라면, GPU를 잘 활용하도록 배치 크기를 4로 유지해야 해요.
| batch size | gradient accumulation steps | effective batch size | |
|---|---|---|---|
| 1 | 64 | 64 | 👎 |
| 4 | 16 | 64 | 👍 |
Gradient checkpointing
그래디언트 체크포인팅(gradient checkpointing)은 backward pass 중 일부 중간 활성화만 저장하고 나머지는 다시 계산해서 메모리 사용을 줄입니다. forward pass의 모든 중간 활성화를 저장하지 않아도 되므로 메모리 오버헤드가 큽니다. 다만 훈련 속도가 느려지는(~20%) 대가가 따릅니다.
from transformers import TrainingArguments
args = TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=16,
gradient_checkpointing=True,
)
혼합 정밀도 (Mixed precision)
혼합 정밀도는 일부 계산을 반정밀도(fp16), 일부를 전체 정밀도(fp32)로 수행해서 훈련 속도를 높입니다. 반정밀도 계산은 전체 정밀도보다 계산 비용이 적어 속도를 높이고, 일부를 전체 정밀도로 유지해 정확도를 보존합니다.
혼합 정밀도 훈련에 사용할 수 있는 데이터 타입이 여러 가지예요. 혼합 정밀도 훈련의 주요 장점은 활성화를 fp16으로 저장하는 것입니다.
TrainingArguments에서 fp16()을 설정해 fp16 혼합 정밀도 훈련을 켭니다.
from transformers import TrainingArguments
args = TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=16,
gradient_checkpointing=True,
fp16=True.
)
fp16은 최적화 단계에서 fp16으로 계산된 그래디언트를 fp32로 다시 변환하므로 메모리 최적화가 아니에요. 특히 작은 배치 크기에서는 GPU에 두 버전(fp16, fp32)의 모델이 생기므로 메모리를 더 쓸 수 있습니다.
bf16은 정밀도 일부를 포기하는 대신 훨씬 큰 동적 범위를 얻어, 오버플로·언더플로 오류를 피하는 데 도움이 됩니다. fp16처럼 loss scaling 방법을 추가할 필요 없이 bf16을 쓸 수 있어요. bf16은 NVIDIA Ampere 아키텍처 이상에서 지원됩니다.
from transformers import TrainingArguments
args = TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=16,
gradient_checkpointing=True,
bf16=True,
)
tf32은 NVIDIA Ampere GPU의 모드로, 컨볼루션과 행렬 곱셈 입력을 tf32로 변환합니다. 나머지 저장·연산은 모두 fp32로 유지돼요. 그래서 tf32는 fp32와 같은 범위, fp16과 같은 정밀도, bf16보다 높은 정밀도를 유지할 수 있습니다. tf32를 fp16 또는 bf16 혼합 정밀도 훈련과 조합하면 처리량을 16배 높일 수 있어요.
tf32는 NVIDIA Ampere GPU에서 기본적으로 켜져 있지만, 아래 코드를 fp32 훈련·추론 코드에 추가해 명시적으로 켤 수도 있습니다.
import torch
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
from transformers import TrainingArguments
args = TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=16,
gradient_checkpointing=True,
bf16=True.
tf32=True,
)
옵티마이저
Transformers는 기본적으로 PyTorch의 AdamW (adamw_torch) 옵티마이저를 구현합니다. 하지만 과거 그래디언트의 가중 평균을 저장하므로, 모델 파라미터 수에 비례하는 추가 메모리가 필요해요. 매우 큰 모델을 훈련할 때는 문제가 될 수 있으며, 그런 경우 다른 옵티마이저를 고려해야 합니다. 예를 들어 NVIDIA나 AMD에 Apex가 설치돼 있다면, adamw_apex_fused 옵티마이저가 모든 AdamW 옵티마이저 중 가장 빠른 훈련을 제공합니다.
TrainingArguments에서 optim()을 설정해 옵티마이저를 고릅니다.
from transformers import TrainingArguments
args = TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=16,
gradient_checkpointing=True,
bf16=True,
optim="adamw_bnb_8bit"
)
훈련 시나리오에 따라 고를 수 있는 옵티마이저가 많습니다(OptimizerNames 참고). 예를 들어 Adafactor는 행렬의 각 요소 대신 행·열의 가중 평균을 저장해 메모리를 크게 줄이지만 수렴은 더 느립니다. 또 bitsandbytes의 8-bit AdamW 옵티마이저로 옵티마이저 상태를 양자화할 수 있어요. 옵티마이저 상태를 낮은 정밀도로 저장하고 옵티마이저 단계에서 사용 전에 디양자화합니다.
데이터 프리로딩
데이터 프리로딩은 배치 데이터를 CPU에서 미리 로딩·준비해 GPU가 계속 작업하게 해서, GPU 유휴 시간을 줄이고 활용도를 높입니다. GPU를 항상 작업하게 만드는 방법은 두 가지예요.
- CPU에 고정 메모리(pinned memory)를 할당해 데이터를 저장하고 GPU로 직접 전송한다.
- 데이터를 더 빨리 프리로드하도록 CPU 스레드 또는 워커 수를 늘린다.
TrainingArguments에서 dataloader_pin_memory()와 dataloader_num_workers()로 고정 메모리 할당과 워커 수를 늘립니다.
from transformers import TrainingArguments
args = TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=16,
gradient_checkpointing=True,
bf16=True,
optim="adamw_bnb_8bit",
dataloader_pin_memory=True,
dataloader_num_workers=4,
)
PyTorch
PyTorch는 메모리 요구량을 줄이고 훈련 속도를 높이는 여러 기능을 제공합니다. Transformers에서 몇 줄만 추가하면 켤 수 있는 경우가 많아요.
torch.empty_cache_steps
torch.cuda.empty_cache 함수는 사용하지 않는 캐시 메모리를 해제해 OOM(메모리 부족) 오류를 피하는 데 도움이 되지만, 훈련이 약 10% 느려집니다.
TrainingArguments의 torch_empty_cache_steps()로 일정 훈련 스텝 후 켤 수 있습니다.
from transformers import TrainingArguments
args = TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=16,
gradient_checkpointing=True,
bf16=True,
optim="adamw_bnb_8bit",
dataloader_pin_memory=True,
dataloader_num_workers=4,
torch_empty_cache_steps=4,
)
torch.compile
torch.compile은 PyTorch 코드를 최적화된 커널로 컴파일해 훈련을 크게 빠르게 합니다. TorchDynamo가 Frame Evaluation API로 PyTorch 그래프를 캡처하고, 이 그래프를 다양한 백엔드용 최적화 커널로 더 컴파일할 수 있습니다.
TrainingArguments에서 torch_compile()을 켜고, torch_compile_backend()로 백엔드를 고릅니다.
from transformers import TrainingArguments
args = TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=16,
gradient_checkpointing=True,
bf16=True,
optim="adamw_bnb_8bit",
dataloader_pin_memory=True,
dataloader_num_workers=4,
torch_empty_cache_steps=4,
torch_compile=True,
torch_compile_backend="inductor"
)
아래 표로 훈련 시나리오에 맞는 백엔드를 고르세요.
| backend | 설명 | 목적 |
|---|---|---|
| eager | PyTorch로 추출한 GraphModule 실행 | 디버깅 |
| aot_eager | AOTAutograd가 추출한 forward·backward 그래프에 PyTorch eager 모드 사용 | 디버깅 |
| inductor | Triton 커널로 TorchInductor + AOTAutograd + CUDA Graphs 사용 | 훈련과 추론 |
| nvfuser | TorchScript와 함께 nvFuser 사용 | 훈련과 추론 |
| aot_nvfuser | AOTAutograd와 함께 nvFuser 사용 | 훈련과 추론 |
| aot_cudagraphs | AOTAutograd와 함께 CUDA Graphs 사용 | 훈련과 추론 |
| ofi | TorchScript의 optimize_for_inference 사용 | 추론 |
| fx2trt | Torch-TensorRT 사용 | 추론 |
| onnxrt | CPU·GPU 추론에 ONNX-RT 사용 | 추론 |
Scaled dot product attention
torch.nn.functional.scaled_dot_product_attention (SDPA)은 스케일드 닷 프로덕트 어텐션 메커니즘의 네이티브 PyTorch 구현이에요. SDPA는 트랜스포머 모델의 원래 어텐션 메커니즘보다 더 효율적이고 최적화되어 있습니다. 세 가지 스케일드 닷 프로덕트 어텐션을 지원해요.
- FlashAttention2은 fp16·bf16 torch 타입 모델에서 자동으로 켜집니다. 먼저 모델을 적절한 타입으로 캐스팅하세요.
- xFormers 또는 Memory-Efficient Attention은 fp32 torch 타입 모델을 지원합니다.
- C++ 구현의 스케일드 닷 프로덕트 어텐션.
SDPA는 PyTorch 2.1.1+에서 기본으로 켜지지만, from_pretrained()에서 attn_implementation="sdpa"로 명시적으로 켤 수도 있습니다.
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B", device_map="auto", attn_implementation="sdpa")
Liger Kernel
Liger Kernel은 RMSNorm, RoPE, SwiGLU, CrossEntropy, FusedLinearCrossEntropy 등 LLM 훈련용 레이어를 단일 Triton 커널로 융합한 모음이에요. 이 커널들은 FlashAttention, FSDP, DeepSpeed와 호환됩니다. 결과적으로 Liger Kernel은 멀티 GPU 훈련 처리량을 높이고 메모리 사용을 줄입니다. 멀티 헤드 훈련과 더 큰 어휘 크기, 더 큰 배치 크기, 더 긴 컨텍스트 길이를 지원하는 데 유용해요.
pip install liger-kernel
TrainingArguments에서 use_liger_kernel=True로 설정해 훈련에 Liger Kernel을 켭니다. 이렇게 하면 모델의 해당 레이어들이 Liger 커널로 패치됩니다.
[!TIP] Liger Kernel은 Llama, Gemma, Mistral, Mixtral 모델을 지원합니다. 최신 지원 모델 목록은 patching 목록을 참고하세요.
from transformers import TrainingArguments
training_args = TrainingArguments(
...,
use_liger_kernel=True
)
liger_kernel_config 파라미터로 적용할 구체적 커널을 설정할 수도 있어요. 이 dict는 _apply_liger_kernel_to_instance 함수에 keyword 인자로 전달되어 커널 사용을 세밀하게 제어합니다. 사용 가능한 옵션은 모델에 따라 다르지만 보통 rope, swiglu, cross_entropy, fused_linear_cross_entropy, rms_norm 등을 포함합니다.
from transformers import TrainingArguments
# Apply only specific kernels
training_args = TrainingArguments(
...,
use_liger_kernel=True,
liger_kernel_config={
"rope": True,
"cross_entropy": True,
"rms_norm": False, # Don't apply Liger's RMSNorm kernel
"swiglu": True,
}
)
NEFTune
NEFTune는 훈련 중 임베딩 벡터에 노이즈를 추가해 모델 성능을 개선합니다. Trainer에서 TrainingArguments의 neftune_noise_alpha 파라미터로 추가되는 노이즈 양을 제어하면서 켤 수 있어요.
from transformers import TrainingArguments, Trainer
training_args = TrainingArguments(..., neftune_noise_alpha=0.1)
trainer = Trainer(..., args=training_args)
원래 임베딩 레이어는 예상치 못한 동작을 피하기 위해 훈련 후 복원됩니다.
더 알아보기 (Learn more)
- 멀티 GPU 훈련: perf_train_gpu_many
- 하드웨어·소프트웨어 고려사항: Performance and Scalability
- 옵티마이저 상세: Optimizer 문서