[Ultralytics] MPS 디바이스에서 메모리 누수 방지를 위한 적극적 메모리 정리

PR 링크: ultralytics/ultralytics#24038 상태: Merged | 변경: +2 / -1

들어가며

Apple Silicon(M1/M2/M3/M4)에서 PyTorch의 MPS(Metal Performance Shaders) 백엔드를 사용한 학습은 편리하지만, 메모리 관리에 고유한 문제가 있습니다. CUDA와 달리 MPS는 메모리 사용량을 fraction으로 조회하는 API가 제한적이며, 메모리 해제가 지연되는 leak 현상이 보고되고 있습니다. 이번 PR은 MPS 디바이스에서 _clear_memory 호출 시 threshold 검사를 건너뛰고 항상 메모리를 정리하도록 합니다.

핵심 코드 분석

validation 전후 메모리 정리

Before:

# Validation
if self.args.val or final_epoch or self.stopper.possible_stop or self.stop:
    self._clear_memory(threshold=0.5)  # prevent VRAM spike

After:

# Validation
if self.args.val or final_epoch or self.stopper.possible_stop or self.stop:
    self._clear_memory(None if self.device.type == "mps" else 0.5)  # prevent VRAM spike

epoch 종료 시 메모리 정리

Before:

self._clear_memory(0.5)  # clear if memory utilization > 50%

After:

# clear if memory utilization > 50%; always clear on MPS due to leak
self._clear_memory(None if self.device.type == "mps" else 0.5)

_clear_memory(threshold) 메서드는 threshold가 None이면 무조건 메모리를 정리하고, 숫자가 주어지면 해당 비율 이상 사용 중일 때만 정리합니다. MPS에서는 메모리 사용률을 정확히 조회하기 어려워 threshold 기반 판단이 신뢰할 수 없으므로, 항상 정리하는 것이 안전합니다.

왜 이게 좋은가

1. MPS 메모리 누수 완화: Issue #22621에서 보고된 MPS 메모리 누수 문제를 직접적으로 해결합니다.

2. CUDA에는 영향 없음: 조건문으로 MPS만 대상으로 하므로 CUDA 사용자의 기존 동작은 변경되지 않습니다. CUDA에서는 threshold 기반의 효율적 메모리 관리가 유지됩니다.

3. 최소 침습적 수정: 2줄의 변경으로 MPS 특유의 문제를 해결하며, 코드 구조를 변경하지 않습니다.

정리

MPS 디바이스에서 _clear_memory의 threshold를 None으로 설정하여 매 validation과 epoch 종료 시 무조건 메모리를 정리하는 간단하지만 효과적인 수정입니다. Apple Silicon에서의 학습 안정성을 높이는 실용적 패치입니다.

참고 자료

• ultralytics/ultralytics#24038
• 관련 Issue #22621

---

이 글은 AI(Claude)의 도움을 받아 작성되었으며, 실제 PR의 코드 변경 사항을 기반으로 분석한 내용입니다.