PyTorch 2.0 정식 출시: torch.compile 도입으로 모델 학습·추론 대폭 가속

PyTorch 2.0은 널리 쓰이는 딥러닝 프레임워크 PyTorch의 차세대 메이저 업데이트입니다. 파이썬 특유의 직관적인 동적 그래프 개발 환경을 그대로 유지하면서도, 내부 컴파일러 계층을 완전히 재설계해 기존 코드를 고치지 않고도 모델 학습과 추론 속도를 끌어올릴 수 있습니다.

이번 버전은 1.13.1 이후 4,541건의 커밋과 428명의 기여자가 참여해 완성됐으며, 기존 코드와의 100% 하위 호환성을 유지합니다.

무엇을 할 수 있고 어떻게 동작하나

핵심 기능은 선택적으로 사용할 수 있는 torch.compile API입니다. 기존 모델 객체를 이 함수로 감싸기만 하면 최적화된 컴파일 모델을 반환받을 수 있습니다.

내부적으로는 4대 핵심 기술이 작동합니다. TorchDynamo가 파이썬 실행 프레임을 분석해 계산 그래프를 안전하게 추출하고, AOTAutograd가 사전 역전파를 추적하며, PrimTorch가 2,000여 개 연산자를 250여 개 기본 연산자로 단순화합니다. 마지막으로 TorchInductor가 하드웨어에 맞게 코드를 생성합니다.

Nvidia 및 AMD GPU 환경에서는 OpenAI Triton 컴파일러를 활용해 숙련자가 직접 짠 CUDA 커널 수준의 코드를 만들어냅니다. 165개 오픈소스 모델 대상 테스트 결과 93%에서 정상 작동했으며, float32 기준 평균 20%, AMP 혼합 정밀도 기준 평균 36%의 속도 향상을 기록했습니다.

트랜스포머 모델을 위한 최적화도 눈에 띕니다. scaled_dot_product_attention 연산자가 torch.nn.functional에 기본 내장되어 FlashAttention이나 xFormers 같은 메모리 최적화 알고리즘을 별도 패키지 설치 없이 즉시 사용할 수 있습니다.

애플 실리콘을 지원하는 MPS 백엔드도 상위 사용 연산자를 포함해 지원 연산자를 300개 이상으로 늘렸고, AWS Graviton3 인스턴스 기반 CPU 추론 성능도 ResNet-50과 BERT 등에서 대폭 개선됐습니다.

시작하는 방법

사용법은 간단합니다. 공식 홈페이지 안내에 따라 pip나 conda로 환경을 업데이트한 뒤 바로 사용할 수 있습니다.

코드에서는 기존 모델 인스턴스 뒤에 model = torch.compile(model) 한 줄만 추가하면 됩니다. 학습 루프나 데이터 로더는 수정할 필요가 없습니다. Mac 환경 사용자라면 최신 macOS로 업데이트해야 개선된 MPS 백엔드의 이점을 온전히 누릴 수 있습니다.

왜 중요한가

대규모 언어 모델을 파인튜닝하거나 최신 딥러닝 파이프라인을 다루는 엔지니어라면 당장 적용해 볼 만합니다. 그동안 번거로운 커스텀 빌드가 필요했던 가속 기능들이 공식 프레임워크 한 줄 코드로 통합되어 도입 장벽이 크게 낮아졌기 때문입니다.

다만 torch.compile은 현재 베타 상태입니다. 벤치마크 모델 대다수를 지원하지만, 매우 복잡한 동적 제어문이나 비표준 사용자 정의 연산자를 사용하는 프로젝트에서는 속도 향상이 제한적일 수 있으므로 프로덕션 적용 전 수치 정합성 검증을 거치는 것이 좋습니다.