PyTorch 2.14, NVGEMM과 장애 복구 강화

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

TL;DR

PyTorch 2.14는 2,995개 커밋과 487명의 기여를 바탕으로 GPU 커널 최적화, 분산 학습 복구, Apple Silicon 연산, 동적 shape 컴파일, 플랫폼 호환성을 넓힌 릴리스입니다. NVGEMM은 CuTeDSL 기반 GEMM 커널에 epilogue fusion과 저정밀 연산을 결합하고, nccl2와 c10d 재구성 인터페이스는 통신 장애가 발생한 대규모 학습 작업을 프로세스 그룹 전체 재시작 없이 복구할 수 있는 경로를 마련합니다. Apple Silicon에서는 Metal 기반 선형대수와 attention 커널이 강화됐으며, MPS의 단일 토큰 decode 경로는 bf16·fp16에서 보고된 8.5배 저하 문제를 보완하고 prefill attention은 기존 커널 대비 약 2~4배 빨라졌습니다. Python 3.15와 3.15t wheel도 추가됐지만 torch.compile은 아직 지원되지 않아 eager 실행만 가능합니다.

섹션별 상세

PyTorch 2.14는 연구용 프레임워크에서 대규모 학습과 추론을 아우르는 hardware-agnostic 플랫폼으로 이동해 온 2.x 계열의 흐름을 이어갑니다. 이번 릴리스에는 2,995개 커밋이 487명의 기여를 거쳐 반영됐으며, 성능·신뢰성·하드웨어 지원이 동시에 확장됐습니다. 특히 NVGEMM, nccl2, fault tolerance, Apple Silicon native linear algebra가 릴리스의 중심 축을 이룹니다.

NVGEMM은 CuTeDSL로 생성한 CUTLASS 커널을 Inductor의 GEMM 후보군에 편입하고 Triton·ATen과 함께 자동 조정합니다. mm, addmm, scaled_mm에서 bias add와 pointwise 연산, GEMM 결과 reduction을 커널 내부에 결합하며 NVFP4의 runtime global scale도 epilogue에서 처리합니다. fused kernel은 디스크에 캐시되고, NVFP4 경로는 Blackwell과 nvidia-cutlass-dsl 4.6.0을 요구합니다.

Inductor는 collective와 독립적인 계산을 교차 배치하는 simple_overlap을 기본 활성화해 통신이 critical path에 남는 시간을 줄입니다. 작은 GPU 커널을 combo kernel로 묶되 큰 reduction은 분리하고, 동적 RBLOCK scaling과 비인라인 sub-kernel로 register pressure를 제어합니다. 이 변경은 별도 설정 없이 분산 학습의 GPU 활용률과 커널 호출 효율을 높이는 방향으로 작동합니다.

bash

# CPU pip3 install torch –index-url https://download.pytorch.org/whl/cpu # CUDA (substitute the CUDA version, e.g. cu126 / cu130) pip3 install torch –index-url https://download.pytorch.org/whl/cu130 # ROCm (substitute the ROCm version) pip3 install torch –index-url https://download.pytorch.org/whl/rocm7.14 # XPU pip3 install torch –index-url https://download.pytorch.org/whl/xpu

Python 3.15와 free-threaded 3.15t용 PyTorch wheel을 CPU, CUDA, ROCm, XPU별 다운로드 인덱스에서 설치하는 명령입니다.

nccl2는 torchcomms에서 포팅된 PyTorch Distributed용 c10d backend로, 재사용 가능한 NcclApi 위에서 full Work contract와 nonblocking communicator를 구현합니다. Backend와 ProcessGroup은 이제 in-place reconfiguration, abort hook, collective 전후 hook을 제공하며 Gloo도 fault-tolerance 경로에 포함됩니다. 한 rank 장애 때 전체 process group을 해체하고 모든 warm state를 버리는 대신 그룹을 현장에서 재구성할 수 있어 대규모 학습 재시작 비용을 줄입니다.

분산 통신은 one-sided RMA window와 backend-agnostic Flight Recorder로 범위를 넓혔습니다. ncclGet과 ncclPut은 peer가 대응 호출을 하지 않아도 메모리를 읽거나 쓸 수 있고, FlightRecorderHook은 ProcessGroup hook을 통해 NCCL뿐 아니라 Gloo와 custom backend의 collective trace를 기록합니다. Python entry point를 통한 backend 등록과 C++·Python 양쪽의 full collective surface 지원도 추가돼 외부 통신 backend 개발 경로가 단순해졌습니다.

DTensor는 전체 device mesh의 모든 placement 조합을 나열하던 규칙을 single-dim strategy 함수로 전환하고 있습니다. matrix·math·tensor 연산이 새 규칙으로 이동하면서 직접 register_op_strategy 등록 수가 158개에서 114개로 줄었고, 등록된 sharding rule을 가진 연산은 2026년 1월 585개에서 1,239개로 늘었습니다. convolution의 일부 마지막 spatial-dimension sharding도 추가돼 조건을 만족하는 경우 forward와 backward에서 all-gather 없이 로컬 처리가 가능합니다.

Apple Silicon에서는 MPSGraph 의존 연산을 손으로 작성한 Metal compute kernel로 옮기고 native linear algebra 범위를 SVD, eigh, lstsq, Cholesky, LU, QR까지 넓혔습니다. SVD·eigh·lstsq는 float32와 complex64에서 Jacobi 계열 커널을 사용하고, Cholesky는 크기에 따라 약 1.2~2.8배 빨라졌으며 lu_factor와 lu_solve는 큰 단일 행렬에서 2~9배, 작은 batched 행렬에서 100배 초과의 측정값이 보고됐습니다. MPS의 단일 토큰 F.linear 경로는 [B, 1, K] 입력을 올바른 GEMV 커널로 보내 bf16·fp16에서 보고된 8.5배 slowdown을 보완합니다.

컴파일과 제어 흐름 측면에서는 torch.switch가 중첩된 torch.cond 대신 index 기반 다중 분기를 제공하고, torch.while_loop는 CUDA conditional node를 사용해 runtime iteration 수를 CUDA graph 안에서 처리합니다. @dynamic_spec은 ShapeVar와 파생 dimension, 제약 조건을 한 번 선언해 torch.compile·torch.export·make_fx에 공유하며, 선언된 dimension은 trace 당시 batch size에 맞춰 조용히 specialize되지 않습니다. complex-valued tensor의 지원 연산은 real·imaginary 계산으로 분해해 compile backend에 전달되지만 아직 모든 complex operation을 지원하지는 않습니다.

플랫폼과 배포 호환성도 바뀌었습니다. Python 3.15와 free-threaded 3.15t용 CPU·CUDA·ROCm·XPU wheel이 제공되고 torchvision 0.29는 torch 2.14와 ABI stable 관계를 갖지만, Python 3.15에서는 torch.compile이 아직 지원되지 않아 eager 실행만 가능합니다. ROCm 7.14 wheel, Intel XPU native graph capture, Rubin sm_107 대상 Inductor, MXFP8·MXFP4와 XPU symmetric memory도 추가돼 AMD·Intel·NVIDIA·Apple 환경의 적용 범위를 넓힙니다.

용어 해설

Inductor: — PyTorch 그래프를 GPU와 CPU에서 실행할 커널 코드로 변환하는 컴파일러 백엔드입니다. 연산을 결합하고 커널 실행 순서를 조정하며 하드웨어별 구현을 선택해 모델 실행의 메모리 접근과 호출 비용을 줄이는 역할을 합니다.
에필로그 융합(Epilogue Fusion): — 행렬 곱셈 결과에 이어지는 bias add, 활성화 함수, 재스케일링 같은 연산을 별도 커널로 실행하지 않고 GEMM 커널 내부에 결합하는 최적화입니다. 중간 결과를 메모리에 기록했다가 다시 읽는 과정을 줄여 메모리 대역폭과 커널 실행 비용을 아낍니다.
RMA 윈도(RMA Window): — 분산 프로세스가 상대 프로세스의 메모리를 직접 읽거나 쓸 수 있도록 제공하는 원격 메모리 접근 인터페이스입니다. 모든 rank가 같은 collective 호출에 참여해야 하는 방식과 달리 필요한 rank만 데이터를 가져오거나 전송할 수 있어 embedding 조회와 expert routing 같은 불규칙한 접근에 적합합니다.
대칭 메모리(Symmetric Memory): — 분산 환경에서 여러 rank가 대응되는 메모리 주소 구조를 공유하도록 할당한 메모리입니다. PyTorch 2.14에서는 CUDA graph 안에서 사용할 수 있고 peer 메모리를 동기화 없이 읽는 get 연산도 제공해 all-gather보다 작은 데이터 이동으로 처리할 수 있는 기반을 마련합니다.
AOTInductor: — PyTorch 모델을 미리 컴파일해 배포 가능한 실행 산출물로 패키징하는 컴파일 경로입니다. 외부 weight 메모리와 pinned 비동기 복사를 활용하면 여러 모델 컨테이너가 같은 GPU weight를 공유하거나 모델 교체 중에도 실행 중인 GPU 작업과 데이터 전송을 겹칠 수 있습니다.

기술

PyTorch 2.14
NVGEMM
CuTeDSL
CUTLASS
Triton
ATen
Inductor
torch.compile
torch.export
make_fx
torch.switch
torch.while_loop
nccl2
torchcomms
Gloo
DTensor
NCCL
NVSHMEM
TokenSwitch
AOTInductor
Helion
MPSGraph
Metal
Metal Performance Primitives
ROCm 7.14
TheRock
Intel XPU
CUDA Graphs
cuBLASLt
Composable Kernel
Origami
FlexAttention
Python 3.15
torchvision 0.29.0

활용 사례

대규모 multi-node 분산 학습에서 node 장애 후 process group을 재구성하는 fault-tolerant training
Mixture-of-Experts 모델의 token dispatch와 expert output combine
Apple Silicon에서 attention, autoregressive decode, SVD, QR, Cholesky를 포함한 학습·추론
가변 batch와 sequence dimension을 사용하는 모델의 torch.compile·torch.export·make_fx 통합
여러 AOTInductor 모델이 동일한 GPU weight를 공유하는 모델 serving
CUDA graph 안에서 symmetric memory와 runtime-dependent while loop를 함께 사용하는 분산 추론