목차
GPU 효율적인 학습
언어 모델을 학습할 때 GPU 메모리는 모델 가중치만 저장하지 않는다. 순전파 결과인 활성값(activation), 역전파의 기울기(gradient), 옵티마이저 상태도 함께 차지한다. 그래서 “모델 파일은 4GB인데 8GB GPU에서 왜 학습이 안 될까?”라는 상황이 생긴다.
이 장에서는 메모리 사용처를 구분하고, 배치·정밀도·체크포인팅을 어떤 순서로 조정할지 살펴본다. 각 기법은 줄이는 메모리와 늘어나는 계산 비용이 다르다.
flowchart TD
A[OOM 또는 학습 속도 저하] --> B{어디서 메모리를 쓰는가?}
B -->|활성값| C[시퀀스 길이·배치 축소]
C --> D[기울기 누적·체크포인팅 검토]
B -->|가중치·옵티마이저| E[혼합 정밀도·LoRA 검토]
E --> F[여러 GPU면 FSDP 검토]
D --> G[품질과 처리량 재측정]
F --> G
1. 메모리 예산부터 계산하기
예를 들어 매개변수가 10억 개인 모델의 가중치를 단순히 16비트로 저장하면 10억 × 2바이트 ≈ 2GB다. 이것은 가중치만의 이론적인 크기다. 학습 중에는 기울기, 옵티마이저 상태, 활성값, CUDA 작업 공간 등이 추가되므로 실제 필요량을 이 숫자로 판단하면 안 된다. 특히 입력 길이와 배치 크기가 커지면 활성값 메모리도 늘어난다.
먼저 실험 기록에 모델 ID·정밀도·최대 토큰 길이·배치 크기·GPU 종류·최대 메모리·초당 처리 토큰을 함께 남긴다. 최적화 한 가지를 적용한 뒤 같은 데이터와 길이로 다시 측정해야 비교가 된다.
2. 배치 크기와 기울기 누적
한 번에 GPU에 올리는 배치를 줄이면 활성값 메모리가 줄어든다. 큰 유효 배치가 필요하다면 여러 미니 배치의 기울기를 모아 한 번 업데이트한다. 아래 예시는 GPU가 1개일 때 2 × 8 = 16개 샘플을 한 업데이트에 반영한다.
from transformers import TrainingArguments
args = TrainingArguments(
output_dir="./checkpoints",
per_device_train_batch_size=2,
gradient_accumulation_steps=8,
)
GPU가 여러 개이면 유효 배치는 대략 GPU 수 × 장치당 배치 × 누적 횟수다. 기울기 누적은 큰 배치를 여러 번에 나눠 계산하는 방법이므로 처리 속도를 높이는 기술은 아니다. 또한 길이가 다른 샘플을 패딩하는 방식과 손실을 토큰 단위로 정규화하는 방식에 따라 같은 유효 배치라도 학습 결과가 조금 달라질 수 있다. Transformers 기울기 누적 문서를 참고한다.
3. 혼합 정밀도와 활성값 다시 계산하기
혼합 정밀도(Mixed Precision)는 연산 일부를 fp16 또는 bf16으로 처리해 메모리와 연산 비용을 줄인다. 하드웨어가 지원하는 형식을 확인하고 선택해야 한다. fp16에서 손실이 NaN이 되는 등 수치 문제가 보이면 손실 스케일과 학습률을 확인하고, 지원 장치라면 bf16을 검토한다.
기울기 체크포인팅(Gradient Checkpointing)은 순전파의 모든 활성값을 저장하지 않고, 역전파 때 일부를 다시 계산한다. 따라서 활성값 메모리는 절약하지만 계산 시간이 늘어난다. Transformers 문서는 이 교환 관계를 설명한다.
args = TrainingArguments(
output_dir="./checkpoints",
per_device_train_batch_size=2,
gradient_accumulation_steps=8,
bf16=True, # GPU의 bf16 지원 여부를 먼저 확인
gradient_checkpointing=True,
)
4. 모델 전체를 갱신할 필요가 있는가
작업이 특정 도메인의 답변 형식 적응이라면 전체 가중치를 갱신하는 대신 LoRA로 작은 어댑터만 학습할 수 있다. 양자화한 기본 모델 위에 LoRA를 학습하는 QLoRA도 선택지다. 다만 양자화가 품질과 연산 지원에 미치는 영향은 모델과 하드웨어에 따라 확인해야 한다. 설정 방법은 PEFT 양자화 가이드에 정리돼 있다.
여러 GPU를 쓸 수 있고 전체 미세 조정이 필요하다면 FSDP(Fully Sharded Data Parallel)로 가중치, 기울기, 옵티마이저 상태를 장치 사이에 나눌 수 있다. 대신 GPU 간 통신과 설정 비용이 추가된다. PyTorch FSDP 문서를 확인한다.
| 증상 | 먼저 확인할 조정 | 치르는 비용 |
|---|---|---|
| 긴 입력에서 OOM | 최대 길이, 패딩, 체크포인팅 | 긴 문맥 축소 또는 재계산 |
| 배치가 커서 OOM | 장치당 배치 축소, 기울기 누적 | 업데이트까지 더 많은 순전파 |
| 가중치·옵티마이저가 큼 | LoRA, 혼합 정밀도, FSDP | 적용 범위 제한 또는 통신 |
| GPU가 바쁘지 않음 | 데이터 로딩, 배치 구성 점검 | 측정과 파이프라인 수정 |
마지막에는 최대 메모리만이 아니라 검증셋 품질과 초당 처리 토큰을 함께 비교한다. 메모리가 절반으로 줄어도 학습이 지나치게 느려지거나 품질이 떨어지면 그 설정은 현재 작업에 맞지 않는다.
실제 학습 배치를 계산해 보기
GPU 2개에서 per_device_train_batch_size=2, gradient_accumulation_steps=8이면 매번 GPU에 올라가는 샘플은 장치당 2개지만, 가중치를 한 번 갱신할 때는 2 × 2 × 8 = 32개 샘플의 기울기가 반영된다. 데이터가 3,200개라면 한 epoch에 대략 100번 업데이트한다. 단, 마지막의 남은 배치, 분산 학습의 샘플 분배, 길이가 다른 문장을 묶는 방식에 따라 실제 업데이트 수는 달라진다. 따라서 max_steps나 스케줄러를 정하기 전에는 학습 로그의 optimizer step 수를 확인한다.
입력 길이도 살펴야 한다. 예를 들어 길이가 2,000 토큰인 문장 하나와 200 토큰 문장 3개를 한 배치에 넣어 모두 2,000 토큰으로 패딩하면, 짧은 문장에도 긴 자리만큼 중간 활성값을 계산해야 한다. 길이가 비슷한 샘플을 묶거나 필요한 문맥 길이까지만 자르는 것이 메모리 개선에 먼저 도움이 될 수 있다. 무조건 길이를 줄이면 정답이 뒤쪽에 있는 학습 사례를 잘라 성능을 잃는다. 잘린 샘플의 비율과 잘린 위치를 기록해야 한다.
메모리와 처리량을 같은 실험에서 기록하기
훈련 스텝을 직접 작성했다면 아래처럼 GPU 최대 사용량을 잴 수 있다. torch.cuda 값은 PyTorch가 관리하는 메모리 통계이며, 다른 프로세스의 사용량이나 GPU 전체 점유량과 같지 않다.
import time
import torch
if not torch.cuda.is_available():
raise RuntimeError("CUDA GPU에서만 이 측정 예시를 실행할 수 있습니다")
torch.cuda.reset_peak_memory_stats()
torch.cuda.synchronize()
started = time.perf_counter()
# 이 위치에서 동일한 학습 데이터로 여러 optimizer step을 수행한다.
torch.cuda.synchronize()
elapsed = time.perf_counter() - started
peak_gib = torch.cuda.max_memory_allocated() / (1024 ** 3)
print(f"경과 시간: {elapsed:.2f}s, PyTorch 최대 할당: {peak_gib:.2f} GiB")
이 코드의 주석 구간을 실제 훈련 루프로 채우고 처음 몇 스텝의 준비 시간은 따로 분리한다. 실험표에는 처리한 토큰 수 / 경과 시간도 남긴다. 아래 수치는 예시다.
| 구성 | 최대 할당 | 토큰/초 | 검증 손실 |
|---|---|---|---|
| 기본 배치 2 | 14 GiB | 1,200 | 1.42 |
| 배치 1 + 누적 2 | 10 GiB | 930 | 1.43 |
| 배치 1 + 누적 2 + 체크포인팅 | 7 GiB | 700 | 1.43 |
이 결과라면 체크포인팅은 GPU 메모리 8GiB 장치에서 학습을 가능하게 할 수 있지만, 16GiB 장치에서는 더 느린 선택이다. 검증 손실의 작은 차이는 우연한 학습 변동인지 반복 실행으로 확인해야 한다. PyTorch CUDA 메모리 문서는 할당량과 캐시된 메모리의 차이를 설명한다.
OOM을 만나면 순서를 정해 진단하기
첫 스텝에서 OOM이면 모델 가중치와 배치의 활성값이 한꺼번에 올라가는 시점을 확인한다. 몇 스텝 뒤에만 발생하면 입력 길이가 긴 배치가 섞였는지, 이전 스텝의 손실 텐서를 리스트에 그대로 보관해 계산 그래프를 붙잡고 있는지 점검한다. loss.item()처럼 숫자로 기록하면 그래프 참조를 피할 수 있다. GPU가 비어 보이는데 느리다면 데이터 로더의 파일 읽기와 토크나이저 처리 시간을 따로 잰다. 원인이 CPU 대기라면 정밀도만 바꾸어도 속도는 개선되지 않는다.