PyTorch 2.15, 분산 학습·저정밀도 추론·다중 GPU 지원 강화

개요

PyTorch 2.15는 2.x 시리즈가 연구 중심 프레임워크에서 대규모 학습과 추론을 위한 생산 플랫폼으로 발전하는 흐름을 이어갑니다. 이번 버전에는 583명의 기여자가 참여한 3,117개 변경 사항이 포함됐습니다. 특정 연산자 하나의 성능보다 분산 통신의 관측성, 메모리 관리, 통신과 계산의 중첩, 하드웨어별 최적화에 폭넓게 초점을 맞췄습니다.

핵심 업데이트

• 재작성된 NCCL 백엔드가 기본값이 됐습니다. nccl을 사용하는 모든 프로세스 그룹은 새 구현을 기본으로 사용합니다. 필요한 NCCL 버전을 갖춘 환경에서는 shrink/grow 경로를 활용해 장애 이후 재구성 과정을 줄일 수 있습니다. 기존 구현으로 되돌리려면 TORCH_DIST_USE_NCCL2=0 또는 nccl-legacy를 사용할 수 있습니다.
• 통신 멈춤 현상 진단이 쉬워졌습니다. 백엔드에 종속되지 않는 플라이트 레코더가 추가되어 collective가 멈췄을 때 각 rank의 사후 정보를 디스크에 기록할 수 있습니다. 한 rank에서 타임아웃이 발생하면 여러 rank의 기록을 함께 확보할 수 있어 대규모 클러스터 장애 분석에 유용합니다.
• 일방향 데이터 전송 API가 추가됐습니다. torch.distributed의 새 Transport API는 첫 백엔드로 NIXL을 지원합니다. 전송을 CUDA stream에 등록할 수 있고 CUDA Graph로 캡처할 수도 있습니다. 프로세스 그룹이나 collective에 의존하지 않는 데이터 이동을 위한 기반입니다.
• FSDP2와 파이프라인 병렬화가 세밀해졌습니다. 정밀도가 중요한 파라미터는 fp32로 유지하면서 나머지는 bf16으로 학습할 수 있습니다. 전체 gradient accumulation을 하나의 CUDA Graph로 캡처할 수 있고, gradient reduction을 다음 계산과 겹칠 수 있습니다. 파이프라인의 수신 버퍼와 전송 한도, unshard 시점도 개선됐습니다.
• 하드웨어별 지원이 확대됐습니다. NVIDIA Blackwell에서는 NVFP4 디코딩을 최적화했고, Apple Silicon에는 FP8, grouped_mm 디코드 경로, 공개 MPS stream API를 추가했습니다. CUDA 13.2가 기본 wheel이 되며 CUDA 13.4와 ROCm 10.0 wheel도 제공됩니다. AMD와 Intel GPU 기능도 확장됐습니다.

의미와 영향

이번 릴리스의 방향은 PyTorch를 단순히 실행 가능한 도구에서 관측하고 조정하며 장애에 대응할 수 있는 기반으로 발전시키는 데 있습니다. 분산 학습 운영팀에는 collective 중단 기록과 FSDP2의 세밀한 제어가 디버깅 및 메모리 관리 부담을 줄이는 요소가 될 수 있습니다. Blackwell과 Apple Silicon을 위한 전용 디코드 경로는 하드웨어별 추론 최적화가 더욱 중요해지고 있음을 보여줍니다.

다만 일부 새 인터페이스는 API Unstable로 표시돼 있습니다. 실제 도입 전에는 드라이버, 가속기, CUDA 또는 ROCm 버전과 기존 분산 학습 구성의 호환성을 검증해야 합니다.

출처: PyTorch Blog