[transformers] Hugging Face Transformers: Apple Silicon(MPS) 환경의 메모리 누수 해결을 위한 MPS Graph Cache 최적화

PR 링크: huggingface/transformers#45818 상태: Merged | 변경: +2 / -0

들어가며

Apple Silicon(MPS, Metal Performance Shaders) 환경에서 딥러닝 모델을 학습할 때, 입력 데이터의 길이가 가변적(variable-length)인 경우 메모리 사용량이 지속적으로 증가하는 현상이 발생합니다. 이는 PyTorch의 MPS 백엔드가 매번 새로운 입력 형태에 대해 그래프를 컴파일하고 이를 MPSGraphCache에 보관하기 때문입니다. 본 PR은 사용자가 기존의 torch_empty_cache_steps 옵션을 활용하여 이 그래프 캐시를 주기적으로 정리할 수 있도록 하여, 장시간 학습 시 발생하는 OOM(Out of Memory) 문제를 효과적으로 제어합니다.

코드 분석

이번 변경사항은 src/transformers/trainer.py 파일의 training_step 메서드 내부에 집중되어 있습니다. 기존에는 clear_device_cache()를 통해 일반적인 디바이스 캐시만 정리했으나, 이제는 MPS 환경일 경우 그래프 캐시까지 함께 정리하도록 로직이 확장되었습니다.

src/transformers/trainer.py

Before:

            if (
                self.args.torch_empty_cache_steps is not None
                and self.state.global_step % self.args.torch_empty_cache_steps == 0
            ):
                clear_device_cache()

After:

            if (
                self.args.torch_empty_cache_steps is not None
                and self.state.global_step % self.args.torch_empty_cache_steps == 0
            ):
                clear_device_cache()
                if torch.backends.mps.is_available() and hasattr(torch.mps, "clear_graph_cache"):
                    torch.mps.clear_graph_cache()

핵심은 torch.mps.clear_graph_cache() 호출을 기존의 캐시 정리 로직과 동일한 조건문 안에 배치한 것입니다. 또한, hasattr을 사용하여 PyTorch 버전이 낮아 해당 API를 지원하지 않는 경우에도 코드가 안전하게 동작하도록 방어적인 프로그래밍을 적용했습니다.

왜 이게 좋은가

성능 및 메모리 효율성

벤치마크 결과에 따르면, 가변 길이 입력이 많은 워크로드에서 기본 정책(always cache)을 사용할 경우 메모리 사용량이 무제한으로 증가합니다. 반면, torch_empty_cache_steps를 통해 그래프 캐시를 주기적으로 정리하면 메모리 증가폭을 크게 낮출 수 있습니다.

• 메모리 절감: 기본 정책 대비 약 450MB 이상의 RSS(Resident Set Size) 증가를 방어할 수 있습니다.
• 유연성: 무조건적인 캐시 삭제는 성능 저하(overhead)를 유발할 수 있습니다. 본 PR은 사용자가 직접 torch_empty_cache_steps를 설정하여 메모리 안정성과 학습 속도 사이의 균형을 맞출 수 있게 합니다.

일반적 교훈

1. 기존 인터페이스 활용: 새로운 옵션을 추가하기보다 기존의 torch_empty_cache_steps를 확장하여 사용자 경험을 일관되게 유지했습니다.
2. 안전한 API 호출: hasattr과 같은 체크를 통해 하위 호환성을 보장하는 것은 라이브러리 개발 시 필수적인 패턴입니다.
3. 환경별 최적화: 특정 하드웨어(MPS)에서만 발생하는 이슈를 범용적인 Trainer 코드에 녹여낼 때는, 해당 하드웨어 사용 여부를 명확히 가드(guard)해야 합니다.

리뷰어 피드백 반영

리뷰 과정에서 벤치마크 스크립트와 테스트 코드를 PR 본문에 포함하지 말라는 피드백을 수용하여, 핵심 로직 변경에만 집중했습니다. 또한, clear_device_cache()와 로직을 통합함으로써 사용자가 별도의 복잡한 설정 없이 기존의 캐시 관리 옵션만으로 MPS 최적화를 누릴 수 있도록 설계되었습니다.

참고 자료

• https://pytorch.org/docs/stable/notes/mps.html
• https://huggingface.co/docs/transformers/main/en/main_classes/trainer#transformers.TrainingArguments.torch_empty_cache_steps

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