PyTorch 2.14, 컴파일러·분산 실행·하드웨어 통합 강화

개요

PyTorch 2.14는 단순히 연산자 목록을 늘린 버전이 아니다. GPU 커널 선택, 분산 통신, 장애 복구, 동적 모델 표현, 여러 가속기 지원을 동시에 다듬어 실제 운영 환경에서 모델을 실행하기 위한 공통 기반을 강화했다. 연구 중심 프레임워크에서 대규모 학습과 추론을 지원하는 통합 플랫폼으로 이동해 온 PyTorch 2.x의 방향이 이번 릴리스에도 이어진다.

핵심 변화

• NVGEMM이 Inductor에 추가됐다. 새 백엔드는 CuTeDSL로 생성한 CUTLASS 커널을 Inductor에 연결한다. Triton 및 ATen과 함께 자동 튜닝 대상이 되며, epilogue fusion, scaled GEMM, NVFP4 GEMM, grouped-reduction epilogue를 지원한다. 행렬 연산에서 컴파일러가 선택할 수 있는 커널 폭이 넓어진 셈이다.
• 분산 통신에 nccl2가 들어왔다. torchcomms에서 포팅된 이 백엔드는 전체 collective contract, 비동기 통신기, 즉시 communicator splitting을 구현한다. 대규모 클러스터에서는 통신 초기화와 재구성이 전체 작업의 효율에 영향을 주므로, 실행 중 유연성이 중요한 개선점이다.
• 장애 허용이 c10d의 핵심 개념이 됐다. 프로세스 그룹을 실행 중에 재구성할 수 있고, 단방향 RMA 윈도우도 사용할 수 있다. Flight Recorder는 NCCL에만 묶이지 않고 모든 백엔드에서 동작한다. 장애 대응과 관측 기능을 특정 통신 구현의 부가 기능이 아니라 공통 런타임 기능으로 다루게 된 것이다.
• Apple Silicon의 선형대수 지원이 확대됐다. Jacobi 커널 기반 SVD, eigh, QR, Cholesky 경로가 추가됐다. 더 많은 MPS 연산과 리덕션이 MPSGraph에서 직접 작성한 Metal 커널로 이동해 그래프 컴파일 비용과 실행 제어의 제약을 줄인다. 단일 토큰 디코딩을 위해 F.linear 경로를 수정하고 GEMV 커널도 추가했다.
• 동적 프로그램을 컴파일하기 쉬워졌다. torch.switch는 torch.cond를 다중 분기로 일반화하며, torch.while_loop는 CUDA Graph에 캡처할 수 있다. @dynamic_spec은 런타임에 변할 수 있는 텐서 차원을 선언적으로 지정하고, 그 정보를 torch.compile, torch.export, make_fx에서 공유한다.
• 하드웨어 범위가 넓어졌다. ROCm 7.14 wheel은 TheRock pip SDK를 통해 생성되고, Intel XPU에는 네이티브 그래프 캡처가 추가됐다. Inductor는 NVIDIA Rubin의 sm_107을 대상으로 한다. 실험적 복소 텐서 컴파일 지원은 일부 연산을 실수부와 허수부 계산으로 분해해 백엔드 최적화를 가능하게 한다.

의미와 영향

모델 개발자 입장에서는 장치별 전용 코드를 작성해야 하는 부담을 줄일 여지가 생겼다. 더 많은 GEMM 커널을 자동으로 선택할 수 있고, 통신과 계산의 오버랩도 기본 활성화되므로 모델 코드를 크게 바꾸지 않고 주요 경로를 개선하기 쉬워진다. 동적 배치, 조건 분기, 반복문을 컴파일과 export 과정에서 일관되게 표현할 수 있다는 점도 실용적이다.

대규모 학습 팀에는 nccl2와 c10d 장애 허용 기능의 결합이 특히 중요하다. 비동기 통신기 관리와 프로세스 그룹 재구성을 통해 노드 장애가 발생했을 때 항상 처음부터 작업을 다시 시작해야 하는 상황을 줄일 수 있다. 백엔드 독립적인 Flight Recorder는 서로 다른 통신 스택의 문제를 진단하는 과정도 단순화할 가능성이 있다.

다만 MPS의 일부 변경은 API 불안정 상태이며, 복소 텐서 컴파일은 실험적 기능이다. 따라서 도입 전에는 실제 모델과 목표 가속기에서 성능과 호환성을 확인해야 한다. 그럼에도 PyTorch 2.14는 성능, 이식성, 신뢰성을 하나의 프레임워크 안에서 함께 개선하려는 전략을 분명히 보여준다.

출처: PyTorch Blog