[ultralytics] YOLO 학습 성능 25% 향상: B200 GPU 최적화 분석

PR 링크: ultralytics/ultralytics#26180 상태: Merged | 변경: +40 / -35

들어가며

최신 딥러닝 모델 학습에서 GPU 연산 효율은 전체 파이프라인의 처리량(throughput)을 결정짓는 핵심 요소입니다. 최근 ultralytics/ultralytics 레포지토리에 병합된 PR은 NVIDIA B200 GPU 환경에서 YOLO 학습 속도를 기존 228.9ms에서 183.5ms로 약 25% 향상시켰습니다. 본 글에서는 이 PR이 해결한 5가지 핵심 병목 현상과 그 기술적 해결책을 분석합니다.

코드 분석

1. Deterministic 모드의 불필요한 메모리 초기화 방지

torch.use_deterministic_algorithms(True)를 사용하면 fill_uninitialized_memory가 활성화되어 모든 신규 할당 텐서를 0으로 채웁니다. 이는 매 스텝마다 수천 개의 커널 호출을 유발합니다.

Before:

# 기존에는 모든 텐서 할당 시 초기화 수행
torch.use_deterministic_algorithms(True)

After:

# Torch 2.2 이상에서 불필요한 NaN fill 커널 비활성화
if TORCH_2_2:
    torch.utils.deterministic.fill_uninitialized_memory = False

2. Fused Adam/AdamW 적용

GradScaler가 호스트에서 found_inf.item()을 읽기 위해 대기하는 시간을 줄이기 위해 fused=True 옵션을 생성자에 전달했습니다.

Before:

optimizer = (partial(MuSGD, muon=muon, sgd=sgd) if use_muon else getattr(optim, name))(params=g)

After:

fused = {"fused": True} if name in {"Adam", "AdamW"} and TORCH_2_0 and next(model.parameters()).is_cuda else {}
optimizer = (partial(MuSGD, muon=muon, sgd=sgd) if use_muon else getattr(optim, name))(params=g, **fused)

3. EMA(Exponential Moving Average) 업데이트 최적화

매 스텝마다 state_dict()를 재구축하던 비효율을 제거하고, 텐서 리스트를 한 번만 빌드하여 재사용하도록 변경했습니다.

Before:

msd = unwrap_model(model).state_dict()
# 매 스텝마다 반복문 수행

After:

if self._pairs is None:
    self._pairs = ema_v, model_v
ema_v, model_v = self._pairs

4. 불필요한 호스트 동기화 제거

max(target_scores.sum(), 1)과 같은 연산은 호스트와 디바이스 간의 동기화를 유발합니다. 이를 clamp_를 사용하여 디바이스 내 연산으로 전환했습니다.

Before:

target_scores_sum = max(target_scores.sum(), 1)

After:

target_scores_sum = target_scores.sum().clamp_(min=1)

왜 이게 좋은가

이 최적화의 핵심 교훈은 'CPU와 GPU 간의 불필요한 대화를 최소화하라'는 것입니다.

1. Stream Sync 감소: 호스트 동기화(Host Sync)는 GPU 파이프라인을 멈추게 합니다. clamp_와 같은 in-place 연산은 이를 방지합니다.
2. 메모리 할당 오버헤드: deterministic 모드에서의 강제 초기화는 비록 안전성을 보장하지만, 학습 성능에는 치명적입니다. 연산이 덮어쓰기(overwrite)를 보장한다면 초기화는 생략 가능합니다.
3. Fused Kernel: fused=True는 여러 커널 호출을 하나로 통합하여 Python 오버헤드와 커널 실행 간의 지연 시간을 획기적으로 줄입니다.

결과적으로 1.25배의 속도 향상을 달성했으며, 이는 대규모 데이터셋 학습 시 수 시간의 단축으로 이어집니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.use_deterministic_algorithms.html
• https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html
• https://pytorch.org/docs/stable/notes/cuda.html#optimizing-cuda-performance

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.