[Hacker News 요약] PyTorch 커스텀 연산: C++/CUDA 구현부터 AOTInductor 컴파일 및 배포까지
28
설명
이 글은 PyTorch 모델에서 C++ 및 CUDA를 활용하여 커스텀 연산(Custom Operation)을 구현하고 사용하는 방법을 상세히 다룹니다. 특히, `torch.export`를 통한 모델 익스포트와 `AOTInductor`를 이용한 컴파일 과정에서 커스텀 연산이 어떻게 통합되는지에 초점을 맞춥니다. 간단한 항등 컨볼루션(identity convolution) 예제를 통해 커스텀 함수와 커스텀 클래스 구현의 전반적인 워크플로우를 제시합니다. 이는 고성능 딥러닝 모델을 개발하고 프로덕션 환경에 배포하려는 개발자에게 매우 유용한 가이드가 될 것입니다.
### 배경 설명
PyTorch는 유연하고 강력한 딥러닝 프레임워크로, 대부분의 모델 개발에 필요한 풍부한 기본 연산들을 제공합니다. 하지만 특정 도메인에 특화된 알고리즘을 구현하거나, 기존 C/C++ 라이브러리와 연동해야 할 때, 또는 극한의 성능 최적화가 필요할 때는 기본 연산만으로는 한계가 있습니다. 이러한 경우 개발자는 C++ 및 CUDA를 사용하여 PyTorch 커스텀 연산(Custom Operation)을 직접 구현하게 됩니다. 이는 모델의 특정 부분을 GPU 하드웨어에 최적화된 코드로 직접 작성하여 성능 병목 현상을 해결하거나, PyTorch가 기본적으로 제공하지 않는 복잡한 로직을 통합하는 데 필수적입니다.
최근 PyTorch는 `torch.compile` 및 `torch.export`와 같은 기능을 통해 모델 배포 및 최적화 워크플로우를 강화하고 있습니다. 특히 `AOTInductor`와 같은 AOT(Ahead-Of-Time) 컴파일러는 모델을 프로덕션 환경에 배포할 때 Python 런타임 의존성을 줄이고 C++ 기반의 고성능 추론을 가능하게 합니다. 이러한 환경에서 커스텀 연산이 제대로 작동하고 최적화될 수 있도록 통합하는 것은 매우 중요하며, 본 글은 이 복잡한 과정을 명확하게 설명하고 있어 개발자들에게 큰 주목을 받고 있습니다. 커스텀 연산의 구현부터 컴파일된 모델에서의 활용까지 전 과정을 다루는 실용적인 가이드라인을 제시한다는 점에서 그 가치가 높습니다.
### PyTorch 커스텀 함수 구현
PyTorch 커스텀 함수는 C++ 및 CUDA로 구현될 수 있으며, `TORCH_LIBRARY_IMPL` 매크로를 사용하여 PyTorch에 등록됩니다. 이 매크로는 CPU와 CUDA 구현을 모두 제공할 수 있게 하여, 입력 텐서의 장치(device)에 따라 PyTorch가 적절한 구현을 자동으로 디스패치하도록 합니다. 예제에서는 `identity_conv_cpu_impl`과 `identity_conv_cuda_impl`을 통해 CPU와 CUDA에서 각각 항등 컨볼루션을 수행하는 방법을 보여주며, 특히 CUDA 커널 구현 시 `AT_DISPATCH_FLOATING_TYPES_AND2` 매크로를 사용하여 다양한 부동 소수점 타입에 대응하는 유연성을 확보합니다. 커스텀 함수는 상태를 가지지 않는(stateless) 연산에 적합합니다.
### PyTorch 커스텀 클래스 구현
매개변수를 유지하고 `forward()` 메서드를 통해 Python에서 호출될 수 있는 상태 저장(stateful) 커스텀 연산이 필요한 경우, `torch::CustomClassHolder`를 상속받아 C++에서 커스텀 클래스를 정의하고 `TORCH_LIBRARY` 매크로로 등록할 수 있습니다. 예제 `IdentityConvClass`는 `channels_`라는 매개변수를 가지며, `forward` 메서드는 내부적으로 앞서 구현한 CUDA/CPU 항등 컨볼루션 구현을 호출합니다. 이 클래스는 `torch.export`와 `AOTInductor` 컴파일을 위해 `__obj_flatten__` 및 `def_pickle`과 같은 추가적인 메타데이터 등록이 필요하며, 이는 직렬화 및 추적(tracing) 호환성을 보장합니다.
### Python에서 커스텀 연산 활용 및 FakeTensor 등록
C++로 구현된 커스텀 함수와 클래스는 공유 라이브러리(`libidentity_conv_ops.so`)로 빌드된 후, Python에서 `torch.ops.load_library`를 통해 로드됩니다. `torch.compile` 및 `torch.export` 호환성을 위해, `FakeTensor` 기반의 심볼릭 트레이싱이 실제 C++/CUDA 코드를 실행하지 않고도 올바르게 작동하도록 '가짜(abstract)' 버전을 등록하는 것이 중요합니다. 이는 `@register_fake_class` 데코레이터를 사용하여 커스텀 클래스의 추상 구현을 제공하고, `@torch.library.register_fake`를 사용하여 커스텀 함수의 추상 구현을 등록함으로써 이루어집니다. 이 과정은 `torch.export`가 모델의 구조와 데이터 흐름을 정확히 파악하는 데 필수적입니다.
### PyTorch 모델 내 커스텀 연산 통합 및 AOTInductor 컴파일
위에서 구현하고 등록한 커스텀 함수와 클래스는 일반적인 `torch.nn.Module` 내에서 다른 PyTorch 연산과 함께 사용될 수 있습니다. 예제 `IdentityModel`은 네 개의 레이어로 구성되며, 이 중 두 개는 네이티브 PyTorch 연산을, 다른 두 개는 커스텀 클래스(`IdentityConvCustomClass`)와 커스텀 연산(`IdentityConvCustomOp`)을 사용합니다. 이렇게 구성된 모델은 `torch.export.export()`를 통해 익스포트될 수 있으며, 익스포트된 그래프는 커스텀 연산이 `torch.ops.higher_order.call_torchbind` 및 `torch.ops.my_ops.identity_conv_op.default`로 표현됨을 보여줍니다. 최종적으로 `torch._inductor.aoti_compile_and_package`를 사용하여 모델을 컴파일하고 `model.pt2` 패키지로 생성하여 배포 준비를 마칩니다.
### Python 및 C++ 추론 환경에서의 활용
컴파일된 `model.pt2` 패키지는 Python 환경에서 `torch._inductor.aoti_load_package`를 통해 로드하여 추론을 실행할 수 있습니다. 이때 커스텀 연산이 포함된 공유 라이브러리가 미리 로드되어 있어야 합니다. 더욱 중요한 점은, 이 패키지가 순수 C++ 추론 프로그램에서도 활용될 수 있다는 것입니다. `dlopen`을 사용하여 커스텀 연산 공유 라이브러리를 로드한 후, `torch::inductor::AOTIModelPackageLoader`를 통해 컴파일된 모델을 로드하여 Python 의존성 없이 고성능 C++ 추론을 수행할 수 있습니다. 이는 임베디드 시스템이나 저지연(low-latency) 서버 환경과 같이 Python 런타임이 부담이 되는 환경에서 특히 유용합니다.
### 가치와 인사이트
이 글은 PyTorch 개발자가 모델의 성능을 극대화하고 특정 하드웨어에 최적화된 로직을 통합하는 데 필요한 핵심적인 기술을 제공합니다. C++/CUDA 커스텀 연산을 통해 PyTorch의 유연성을 유지하면서도 네이티브 코드 수준의 성능을 달성할 수 있습니다. 특히 `torch.export`와 `AOTInductor`를 활용한 컴파일 및 배포 워크플로우는 모델을 연구 단계에서 프로덕션 환경으로 전환하는 데 있어 중요한 가교 역할을 합니다. Python 환경뿐만 아니라 순수 C++ 환경에서도 컴파일된 모델을 실행할 수 있게 함으로써, 배포의 유연성과 효율성을 크게 향상시킵니다. 이는 딥러닝 모델의 실제 서비스 적용 가능성을 넓히고, 개발자가 성능과 배포 용이성이라는 두 마리 토끼를 모두 잡을 수 있도록 돕는 실무적인 통찰을 제공합니다.
### 기술·메타
- PyTorch
- C++
- CUDA
- AOTInductor
- torch.export
- FakeTensor
- dlopen
### 향후 전망
PyTorch 커스텀 연산의 구현 및 배포는 앞으로도 중요한 개발 영역으로 남을 것입니다. 향후 PyTorch는 `torch.compile` 및 `AOTInductor`와 같은 컴파일러 스택을 더욱 고도화하여, 커스텀 연산의 통합 과정을 더욱 간소화하고 자동화하는 방향으로 발전할 가능성이 높습니다. 예를 들어, 커스텀 연산의 `FakeTensor` 등록이나 `__obj_flatten__`과 같은 메타데이터 정의를 위한 고수준 API나 도구가 제공되어 개발자의 부담을 줄일 수 있습니다. 또한, 다양한 하드웨어 가속기(NPU, FPGA 등)에 대한 지원이 확대됨에 따라, 커스텀 연산을 통해 이러한 새로운 하드웨어의 잠재력을 최대한 활용하는 방법론이 더욱 중요해질 것입니다.
경쟁 측면에서는 TensorFlow, JAX 등 다른 딥러닝 프레임워크들도 유사한 커스텀 연산 및 컴파일러 통합 기능을 제공하고 있어, PyTorch는 개발자 경험과 성능 면에서 지속적인 우위를 점하기 위해 노력할 것입니다. 커뮤니티는 더 많은 최적화된 커스텀 연산 라이브러리를 공유하고, 복잡한 모델을 효율적으로 배포하는 모범 사례를 발전시킬 것으로 예상됩니다. 궁극적으로는 연구 단계에서 구현된 혁신적인 아이디어가 프로덕션 환경에서 지연 없이, 그리고 최적의 성능으로 구동될 수 있도록 하는 것이 목표가 될 것입니다.
📝 원문 및 참고
- Source: Hacker News
- 토론(HN): [news.ycombinator.com](https://news.ycombinator.com/item?id=48358339)
- 원문: [링크 열기](https://leimao.github.io/blog/PyTorch-Custom-Operation/)
---
출처: Hacker News · [원문 링크](https://leimao.github.io/blog/PyTorch-Custom-Operation/)
제목글쓴이조회
- [AI Breakfast] OpenAI는 6억 명 이상의 사용자를 대상으로 에이전트 슈퍼앱 전환을 추진합니다. 구글은 2026년 10월부터 SpaceX와 월 9억 2천만 달러 규모의 GPU 컴퓨팅 계약을 체결했습니다. Anthropic은 2025년 후반 모델을 능가하는 미공개 모델과 맞춤형 칩 개발에 집중하고 있습니다.[0]Nedai12
- [The Verge] 마이크로소프트 AI 총괄, '초지능 임박했지만 일자리는 빼앗지 않아' 선언[0]Nedai12
- [The Verge] 구글 NotebookLM, 제미니 3.5 업그레이드로 클라우드 기반 연구 역량 강화[0]Nedai13
댓글 0
아직 댓글이 없습니다. 첫 댓글을 남겨 보세요.