PyTorch 2.15 MPS 릴리스 노트

개요

PyTorch 2.15는 네이티브 FP8 지원, 최적화된 LLM 디코딩 경로, 전용 Metal 커널로의 전환을 도입하여 Apple Silicon 사용자를 위한 성능 및 메모리 효율성을 크게 향상시켰습니다. 이러한 업데이트는 MPSGraph에 대한 의존도를 줄여 밀집(dense) 모델 및 전문가 혼합(MoE) 모델의 추론 속도를 높이고, 비디오 생성과 같은 복잡한 워크로드에서 메모리 사용량을 더욱 안정적으로 만듭니다.

주요 기술 개선 사항

FP8 지원

PyTorch 2.15는 MPS 백엔드에서 torch.float8_e4m3fn 텐서에 대한 지원을 추가합니다. 이를 통해 개발자는 요소당 1바이트만 사용하여 값을 저장할 수 있으므로 FP16 또는 BF16 대비 메모리 사용량을 효과적으로 절반으로 줄일 수 있습니다. 이번 릴리스에는 GPU에서 직접 FP8 행렬 곱셈을 수행하기 위한 F.scaled_mm이 포함되어 있으며, 텐서 단위 및 행 단위 스케일링을 모두 지원합니다.

LLM 디코딩 최적화

MPSGraph를 우회하고 네이티브 Metal 커널을 사용함으로써 밀집 모델과 MoE 모델 모두의 디코딩 성능이 크게 향상되었습니다:

  • Linear Layers: F.linear는 이제 특수 행렬-벡터 커널을 활용하여 일반적인 디코딩 형태에서 기하 평균 1.23배의 속도 향상을 제공합니다.
  • MoE 모델: 새로운 네이티브 그룹화된 행렬 곱셈(grouped matmul) 구현으로 더 빠른 토큰 라우팅이 가능해졌습니다. torch.compile을 사용하는 M3 Ultra에서 Qwen3.6-35B-A3B와 같은 MoE 모델의 처리량은 초당 9.4 토큰에서 61.3 토큰으로 증가했습니다.
  • Softmax: 전용 Metal 커널을 통해 softmax 및 log_softmax 연산의 지연 시간이 이전 MPSGraph 구현 대비 약 8배 감소했습니다.

비디오 워크로드의 메모리 관리

MPSGraph 기반 패딩을 전용 Metal 커널로 대체함으로써 PyTorch 2.15는 비디오 처리 시 발생하는 상당한 메모리 오버헤드 문제를 해결했습니다. 예를 들어, Wan2.2를 위한 타일형 VAE 디코딩 프로세스는 이전 버전의 270 GiB에서 8 GiB로 메모리 사용량이 감소했습니다.

선형 대수 및 연산

선형 대수 기능이 다음과 같이 확장되었습니다:

  • 복소수 LU 분해, LU 풀이 및 삼각 행렬 풀이.
  • 소형 행렬 역행렬(최대 4x4)을 위한 최적화된 커널 및 SVD 연산의 견고성 향상.
  • conv1d, 상수/반사/복제 패딩 및 업샘플링 역방향을 위한 새로운 Metal 기반 경로.

개발자 제어 및 요구 사항

개발자는 이제 torch.mps.Stream 및 torch.mps.stream 컨텍스트 관리자를 통해 실행을 더 세밀하게 제어할 수 있으며, 이를 통해 명시적인 스트림 관리 및 동기화가 가능합니다.

시스템 요구 사항:

  • macOS: 15 이상.
  • Python: 3.11 이상.

향후 전망: PyTorch 2.16

2.16 릴리스 개발은 MPSGraph 의존도를 더욱 줄이는 데 중점을 두고 있습니다. 계획된 개선 사항은 다음과 같습니다:

  • Full Metal Attention: 스케일드 점곱 어텐션(scaled dot-product attention)을 완전히 Metal 커널로 전환.
  • 명령 버퍼 최적화: 디코딩 중 동시 GPU/CPU 실행을 허용하기 위해 16개의 커널마다 명령 버퍼를 커밋.
  • 바이너리 크기 축소: libtorch_cpu.dylib의 크기를 줄이기 위해 Metal 3.1용 번들 셰이더를 기본적으로 컴파일.

Sources