PyTorch 2.15: 통신 기본값 전환, 플랫폼 성능 개선

1. GPU 사이 통신을 새 구현으로 전환한다

엔비디아 GPU 사이에서 데이터를 주고받는 통신 계층인 NCCL 백엔드가 새 구현으로 바뀐다. nccl을 지정하면 새 구현이 선택되며, 통신 그룹을 만들 때 GPU 사이 통신 채널도 미리 만든다.

통신이 멈췄을 때 원인을 찾는 기록 기능도 더해졌다. 한 학습 프로세스에서 시간 초과가 발생하면 전체 프로세스가 당시 작업 상태를 디스크에 남겨, 멈춤 시점에 무엇을 하고 있었는지 살필 수 있다.

아래는 공식 릴리스 노트의 초기화 예시다. 기본 구현 외에 상대별로 독립된 통신 채널을 쓰는 방식과 기존 구현을 선택하는 방식도 제시한다.

dist.init_process_group("nccl", device_id=torch.device("cuda", local_rank))  # new implementation
dist.init_process_group("nccl-lazy", device_id=torch.device("cuda", local_rank))  # independent P2P communicators
dist.init_process_group("nccl-legacy")  # legacy implementation, removed in the next release

2. 학습 정밀도와 메모리 사용 시점을 조절한다

모델 상태를 여러 GPU에 나눠 보관하는 FSDP2는 정밀도 제어를 세분화했다. 라우터 게이트나 정규화 가중치처럼 민감한 파라미터는 32비트로 유지하면서 나머지 모델은 낮은 정밀도로 학습할 수 있다.

reduce_dtype=None이면 기울기는 각 파라미터의 기울기 정밀도 설정을 따라 누적·축소된다. 공식 릴리스 노트의 32비트 파라미터·bfloat16 모델 계산 사례에서는 기울기 처리도 16비트에서 32비트로 바뀌어, 기울기를 나눠 합치는 통신의 데이터 양이 2배가 된다. 높은 기울기 정밀도를 유지하는 대신 통신량을 더 쓰는 변화로, 이전 동작을 유지하려면 축소 정밀도를 직접 지정하면 된다.

모델을 여러 단계로 나눠 실행하는 파이프라인 병렬화는 수신 직전에 버퍼를 할당하고 전송이 끝난 출력을 해제한다. 버퍼를 미리 쌓아두는 대신 필요한 시점에 확보해 최대 메모리 사용을 줄이는 방식이다. 기울기 축소와 후속 계산을 겹치는 조정 기능도 추가됐다.

분산 통신에는 통신 그룹이나 집합 통신 없이 저장소를 통해 상대와 연결하는 단측 전송 API가 더해졌다. 첫 백엔드는 NIXL이며, 전송을 GPU 실행 흐름에 올려 일련의 작업을 기록·재실행하는 CUDA 그래프에 담을 수 있다.

3. 애플 실리콘에서 문장 생성 연산을 가속한다

애플 실리콘의 GPU 실행 계층인 MPS는 더 많은 연산을 네이티브 Metal 커널로 옮겼다. 중간 실행 계층을 거치던 연산을 직접 처리하는 커널로 바꿔 호출 부담을 줄이는 방향이다.

공식 발표 기준, 한 행에서 128,256개 어휘의 점수를 확률로 바꾸는 softmax의 소요 시간은 8분의 1로 줄었다. 크기를 늘리는 업샘플링 연산의 역전파는 최대 21.7배 빨라졌다. 텐서 전체의 크기를 계산하는 노름 연산도 100배 넘게 빨라졌다는 결과를 제시했다.

문장을 생성하는 디코드 단계의 F.linear는 공식 발표 기준 기하평균, 곱셈을 바탕으로 계산한 평균인 1.23배, 작은 투영 연산에서는 최대 5배 빨라졌다. 혼합 전문가 모델처럼 일부 하위 신경망을 골라 쓰는 모델의 디코드에서는, M3 Ultra의 새 그룹화 행렬 연산 경로가 처리량을 1.32~1.34배 높였다.

별도 선형대수 결과로, 공식 발표 기준 512×512 행렬의 QR 분해 시간은 약 7초에서 1.5밀리초로 줄었다. QR 분해는 행렬을 직교 행렬과 삼각 행렬로 나누는 계산이다.

4. 설치 패키지와 실행 환경을 함께 선택한다

기본 설치 패키지인 휠은 CUDA 13.2로 바뀌었고 CUDA 13.4 휠도 제공된다. 두 휠은 580 이상의 엔비디아 드라이버와 튜링 세대 이상 GPU를 요구한다. CUDA 12.6이나 13.0 휠을 유지할 환경에는 공식 릴리스 노트가 이전 버전 사용을 안내한다.

파이썬 지원 범위는 3.11~3.15다. 이전 3.10 환경에서 업그레이드하려면 파이썬도 함께 올려야 한다. 애플 실리콘의 MPS 가속에는 macOS 15 이상이 필요하다.

AMD에는 ROCm 10.0 휠이 추가됐다. ROCm 휠은 TheRock 패키지를 기반으로 바뀌어, 라이브러리를 직접 묶는 대신 의존성으로 설치한다.

5. 기존 학습 구성과 저장 데이터의 전환을 준비한다

새 통신 백엔드는 미사용 그룹에도 GPU 메모리를 할당하고 그룹마다 장치 하나에 묶이므로, 한 프로세스가 여러 GPU를 쓰는 구성은 기존 구현이 필요하며 이 구현과 환경 변수를 통한 복귀 경로는 다음 릴리스에서 제거될 예정이다.

2.15에서 저장한 분산 체크포인트는 2.14 이하에서 읽을 수 없어, 이전 버전으로 되돌릴 계획이 있다면 저장 데이터의 호환성까지 고려해야 한다.

PyTorch 2.15의 도입 가치는 연산별 속도 개선과 학습 제어 확대에 있으며, 전환 계획은 통신 구성과 저장 데이터까지 아울러야 한다.