[TensorRT] PyTorch TensorRT: 메모리 복사 제거를 통한 엔진 빌드 최적화

PR 링크: pytorch/TensorRT#4794 상태: Merged | 변경: +154 / -15

들어가며

PyTorch TensorRT 환경에서 대규모 모델(예: 3 GiB 이상의 가중치)을 컴파일할 때, 호스트 메모리 사용량이 급증하는 문제가 있었습니다. 기존 구현은 torch.classes.tensorrt.Engine을 생성하기 위해 엔진 플랜(plan)을 bytes 리스트로 전달했는데, 이 과정에서 TorchBind가 리스트를 std::vector<std::string>으로 변환하며 데이터를 두 번 복사했습니다. 이로 인해 3 GiB 엔진을 빌드할 때 약 6 GiB의 추가 호스트 메모리가 필요했고, 직렬화된 데이터를 직접 역직렬화하는 것보다 훨씬 느린 1.7초의 시간이 소요되었습니다. 본 PR은 이러한 불필요한 복사를 제거하여 메모리 효율과 속도를 극적으로 개선했습니다.

코드 분석

1. C++ 레이어: std::string_view 도입 (core/runtime/TRTEngine.cpp)

기존에는 std::string을 통해 데이터를 전달받아 복사가 발생했습니다. 이를 std::string_view로 변경하여 원본 버퍼를 직접 참조하도록 수정했습니다.

// Before
TRTEngine::TRTEngine(const std::string& mod_name, const std::string& serialized_engine, ...)

// After
TRTEngine::TRTEngine(const std::string& mod_name, std::string_view serialized_engine, ...)

또한, deserializeCudaEngine 호출 시 .c_str() 대신 .data()를 사용하여 메모리 접근 방식을 최적화했습니다.

2. Python 레이어: create_cpp_engine 유틸리티 (py/torch_tensorrt/dynamo/runtime/_serialized_engine_layout.py)

Python 측에서 엔진 정보를 패킹할 때 torch.frombuffer를 사용하여 메모리 복사 없이 CPU 텐서로 뷰(view)를 생성합니다.

def create_cpp_engine(serialized_info: Sequence[Any]) -> Any:
    # ...
    plan_tensor = torch.frombuffer(plan, dtype=torch.uint8)
    return torch.classes.tensorrt.Engine.from_engine_tensor(info, plan_tensor)

3. TorchBind 등록 (core/runtime/register_jit_hooks.cpp)

from_engine_tensor 정적 메서드를 추가하여, C++ 측에서 직접 텐서의 데이터 포인터를 참조하도록 구현했습니다.

.def_static("from_engine_tensor", [](std::vector<std::string> serialized_info, at::Tensor serialized_engine) -> c10::intrusive_ptr<TRTEngine> {
    return c10::make_intrusive<TRTEngine>(
        serialized_info,
        std::string_view(static_cast<const char*>(serialized_engine.const_data_ptr()), serialized_engine.nbytes()));
})

왜 이게 좋은가

이번 최적화의 핵심은 Zero-copy입니다. 기존에는 엔진 플랜이 파이썬에서 C++로 넘어가는 과정에서 데이터가 여러 번 복제되었으나, 이제는 torch.frombuffer와 std::string_view를 통해 동일한 메모리 영역을 공유합니다.

• 성능 수치: 엔진 설정 시간이 1.7초에서 0.14초로 약 12배 단축되었습니다.
• 메모리 효율: 3 GiB 모델 기준, 컴파일 시 피크 호스트 메모리 사용량이 4.66W에서 2.77W로 감소했습니다.
• 일반적 교훈: 대규모 데이터를 Python과 C++ 간에 주고받을 때, bytes 객체나 std::string을 통한 복사는 치명적인 성능 저하를 유발합니다. torch.frombuffer와 같은 텐서 뷰잉(viewing) 기술을 활용하면 메모리 복사를 방지하고 시스템 자원을 효율적으로 사용할 수 있습니다.

리뷰어 피드백

리뷰 과정에서 serialized_engine이 읽기 전용 버퍼임을 고려하여 warnings.catch_warnings()를 통해 경고를 억제하는 처리가 추가되었습니다. 이는 torch.frombuffer가 기본적으로 쓰기 가능 버퍼를 기대하기 때문에 발생하는 불필요한 경고를 방지하기 위함입니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.frombuffer.html
• https://pytorch.org/docs/stable/cpp_class.html

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.