Apple Silicon에서 Vesuvius Challenge 추론 속도가 9.7배 느렸던 이유를 수정하다

Vesuvius Challenge의 공식 튜토리얼(tutorial5, 잉크 검출 추론)을 Apple Silicon Mac에서 그대로 실행했더니, 추론이 명확하게 느린데도 어떠한 경고도 나오지 않았습니다. 원인을 추적해 보니, 디바이스 선택 코드가 CUDA만 확인하고 있었기 때문에, MPS를 사용할 수 있는 환경에서도 무경고로 CPU에 폴백(fallback)하고 있었습니다. 같은 패턴을 다른 곳에서도 발견하여 6개 파일에 정리하고 수정했으며, 공개된 스캔 데이터와 공개 체크포인트로 전후를 실측했습니다.

무엇이 일어났는가

tutorial5는 villa이라는 리포지토리의 코드로 잉크 검출 추론을 실행하는 내용입니다. MPS가 PyTorch에서 사용 가능한 Apple Silicon Mac으로 실제 Vesuvius Challenge의 세그먼트와 공개된 ink_9um 체크포인트를 사용하여 테스트했습니다. 3600×3600의 클롭(3136 패치)에 걸린 시간은 약 49분이었습니다. 이 하드웨어로서는 너무 느리다는 느낌이 들었습니다. 로그에는 어떤 디바이스가 선택되었는지에 대한 줄이 하나도 없었고, 이를 알아차릴 수 있는 계기는 코드를 읽어보기 전까지는 없었습니다.

CUDA의 유무만으로 분기하는 코드는 드문 것은 아닙니다. CUDA가 없는 환경에서는 CPU 분기로 들어가는 것이 일반적인 구조입니다. 하지만 Apple Silicon Mac은 CUDA가 없어도 MPS라는 선택지를 가지고 있습니다. 이를 확인하지 않는 코드 위에서는 GPU가 없는 환경과 같은 취급을 받게 되어, 본인이 느끼기 전까지 큰 느림함을 조용히 안고 가게 됩니다.

원인

추론의 디바이스 해결이 CUDA의 유무만 보고 있었습니다. CUDA GPU가 없는 Mac에서는 else 분기로 CPU에 떨어지고, 그대로 아무 경고 없이 실행됩니다. autocast도 CUDA 외에서는 유효하지 않게 만들어져 있어, CPU 실행은 디바이스가 느린 것에 더해 fp32 상태로 계산하고 있었습니다. 이 부분을 추적하는 과정에서, 같은 CUDA 제한 패턴을 dice 계열 메트릭 계산과 lasagna CLI의 디바이스 자동 선택에서도 발견했습니다. 총 3곳이었습니다.

수정한 내용

CUDA를 사용할 수 없으면 MPS를 확인하고, 그것도 안 되면 CPU라는 순서로 변경했습니다. 디바이스가 mps로 결정되었을 때는 로그에 남도록 하고, autocast도 device.type이 cuda 또는 mps일 경우 유효하도록 분기를 수정했습니다. 건드린 것은 추론 본체의 ink_detection/inference 하위 3개 파일(디바이스 해결을 모아둔 inference_runtime.py 포함), dice 계열 메트릭 2개 파일, lasagna CLI의 1개 파일을 합쳐 총 6개 파일이며, 차이점은 +28/-9 줄입니다.

실데이터 전후 비교

같은 세그먼트, 같은 클롭, 같은 체크포인트, 같은 실행 시 인수를 사용하여 디바이스만 바꿔서 비교했습니다. 사용한 것은 세그먼트 PHerc0139/w035의 9.362µm surface-volume(공개 S3, 익명 접근, 전체 두루마리에는 적용하지 않음)와 체크포인트 scrollprize/ink_9um의 hybrid_3d2d-seed42/step-075000.pth입니다. 세그먼트 내부에서 잘라낸 3600×3600px 클롭으로, 선택된 패치는 두 조건 모두 같은 3136장이며, --tta-mirror --overlap 0.5 --no-compile --num-workers 0 --batch-size 1 --amp-dtype auto라는 옵션도 동일하게 사용했습니다.

본격적인 클롭과 패치 수를 결정하기 전에, 작은 수량으로 캘리브레이션을 진행했습니다. 9장으로 6초, 121장으로 25초, 121장에 TTA 미러를 추가하면 92초입니다. 여기서 역산하여 3136 패치라는 규모를 결정했지만, 실제로 CPU로 실행한 결과는 캘리브레이션 예측보다 약간 무거워 49분 11초가 걸렸습니다.

출력이 변하지 않았음을 확인하기

장치를 바꿨을 뿐인데 결과까지 달라진다면 수정한 의미가 희박하므로, 두 조건에서 출력된 확률 맵(3600×3600, uint8)을 화소 단위로 비교했습니다. 평균 절대 차이는 0.0326, 중앙값 절대 차이는 0.0, 최대 절대 차는 255단계 중 1, 표준 편차는 0.178이었습니다. 차이가 1을 초과한 화소는 0%였고, 상관 계수는 0.999981이었습니다. CPU 쪽은 autocast가 없는 fp32이고, MPS 쪽은 체크포인트를 지정한 대로 fp16의 autocast라는 조건 차이가 있습니다. 그럼에도 불구하고 이 정도 수준의 차이는 반올림 오차 범위라고 판단합니다. 출력 화소값의 범위와 평균도 두 조건에서 거의 일치했습니다. CPU 쪽은 53부터 210 범위에 평균 83.764였고, MPS 쪽은 53부터 210 범위에 평균 83.792였습니다. 이미지 자체는 여기에나 리포지토리에도 올리지 않았습니다. 게시한 것은 집계된 수치뿐입니다.

실행 환경 관련 확인 사항

환경은 MacBook Pro(Apple M3 Pro, 메모리 18GB, macOS 26.5), Python 3.11.7, torch 2.14.0(mps_available=True, mps_built=True)입니다. 패키지 메타데이터는 requires-python>=3.14,<3.15를 선언하고 있지만, 코드를 읽어본 결과 실행 시 3.14 고유의 구문이나 기능을 사용하지 않았음을 확인했기 때문에 설치하지 않고 PYTHONPATH를 통해 3.11에서 호출했습니다. 모델 구축과 추론 모두 문제없이 작동합니다. villa는 main에 있는 특정 시점(커밋 bcf631d)을 사용했습니다. 로컬 의존성에서는 nnunetv2와 batchgeneratorsv2를 포함된 경로에서 pip install --no-deps -e로 도입했으며, dynamic-network-architectures는 pypi의 최신 버전을 사용합니다. nnunetv2 측이 선언한 버전 상한과 차이가 있었지만, 모델 구축과 추론 어느 쪽에서도 실질적인 문제는 없었습니다. 이번에는 모델 학습은 전혀 진행하지 않고, 추론만을 대상으로 했습니다.

데이터 획득량도 측정했습니다. 메쉬와 surface-volume의 zarr로 총 약 970MiB, 체크포인트가 약 132MiB입니다. 처음에는 더 큰 양을 예상했지만, 이 세그먼트는 튜토리얼이 지정하는 자체 렌더링(volume-cartographer의 C++ 도구로 메쉬에서 surface-volume을 만드는 절차)을 거치지 않아도, 이미 렌더된 zarr가 공개되어 있었습니다. 이는 해당 세그먼트에 특수한 사정이며, 다른 세그먼트에서도 동일하게 적용될 것이라고는 할 수 없습니다.

확인한 것과 아직 확인하지 않은 것

패치를 작성하기 전에 유사한 증상을 보고한 기존 Issue와 PR을 확인했습니다. 둘 다 koine_machines.inference라는 모듈이 대상이었고, MPS 분기가 없고 경고 없이 CPU로 떨어지는 같은 종류의 증상이었지만, 모듈은 달랐습니다. 이번 수정은 ink_detection의 추론, dice 계열 메트릭스, lasagna CLI의 3곳에서 이루어졌으며, 해당 PR이 변경한 파일과는 겹치지 않습니다. 먼저 Issue를 발견했던 측에 통합하는 것도 좋을 것이라고 생각합니다.

아직 확인하지 못한 것은 동일한 CUDA 한정 패턴이 리포지토리의 다른 곳에도 남아있는지 여부입니다. 이번에 찾은 3곳 외에는 찾아보지 않았습니다.

제출은 추후 진행 예정

지금까지의 내용과 패치를 villa의 자신의 포크(fork)로 push한 후 Issue를 생성하고, Progress Prize 월간 제출 양식에 올릴 예정입니다. 제출이 완료되는 대로 이 기사에 포크와 Issue 링크를 추가하겠습니다.

직접 만들어 본 경험을 바탕으로

튜토리얼을 실행하다가 '느리다'고 느낀 부분에서 시작된 패치입니다. 처음부터 벤치마킹을 할 생각으로 손을 움직인 것은 아니었고, 기다리는 동안 로그를 보다가 장치 선택 코드에 도달했습니다. 실제로 자신의 Mac에서 느림을 경험하고 고쳤기 때문에, 수정한 후의 속도 역시 같은 Mac에서 확인할 수 있었던 점이 좋았습니다.

속도만 보여주고 끝내지 않기 위해, 출력값이 수치적으로 얼마나 가까운지도 함께 보여주기로 했습니다. 장치를 바꾸면 결과까지 달라져서 수정한 의미가 희박해질 수 있기 때문입니다. 화소 단위로 비교하여 반올림 오차 범위 내에 들어와 있음을 확인한 후에 제출 준비를 진행했습니다.

이번에 수정한 곳은 발견된 3군데뿐입니다. 같은 패턴이 다른 곳에도 남아있을지 모르니, 그곳을 다음번에 찾아가고 싶습니다.

논의 (Discussion)