PyTorch가 Dynamic Tensor Shape를 처리하는 법

AI 모델은 프롬프트 길이, 이미지 배치, 영상 클립처럼 입력 크기가 계속 달라지는 환경을 다뤄야 한다. 이런 dynamic tensor shape는 성능 최적화의 큰 장애물이지만, PyTorch 2의 torch.compile은 이를 실용적으로 처리한다.

PyTorch의 기본 실행 방식인 eager mode는 한 번에 하나씩 연산을 수행한다. GPU나 Furiosa의 RNGD 같은 accelerator에서는 Python과 하드웨어 사이의 왕복 오버헤드가 커지므로, PyTorch는 대부분의 연산을 비동기로 실행해 그 비용을 숨긴다.

하지만 오버헤드를 숨기는 것과 없애는 것은 다르다. 그래서 PyTorch 2는 프로그램을 더 compiler-friendly한 dataflow graph로 바꿔, 더 큰 덩어리로 최적화하고 실행하려고 한다. 이를 위해 Python bytecode를 읽고, 호출된 함수들을 따라가며 모델을 추적한다.

이 과정에서 fake tensors가 핵심 역할을 한다. 실제 데이터를 담지 않은 텐서로 코드를 실행하면서 연산과 tensor shape를 기록해, 메모리를 쓰지 않고도 FX graph를 만든다. 예시처럼 x * 2 + 3은 f32[8] 같은 구체적 shape를 가진 FX graph로 바뀌고, 이후 최적화나 kernel 생성의 입력이 된다.

문제는 첫 번째 컴파일이 보통 static shape를 전제로 이뤄진다는 점이다. 이후 다른 shape가 들어오면 같은 코드가 안전하게 실행되는지 확인해야 하므로, PyTorch는 다시 컴파일하면서 해당 텐서를 dynamic shape로 표시한다. 나머지 텐서는 여전히 static으로 둘 수 있어, 필요한 부분만 유연하게 만든다.

이때 PyTorch는 shape 제약을 symbolic하게 추적한다. 예를 들어 덧셈은 두 텐서의 shape가 같거나 broadcast 가능해야 하고, matmul은 contracting dimension이 같아야 한다. 이런 조건을 바탕으로 guard를 만들며, SymPy로 논리식을 줄인 뒤 C++ 검사 함수까지 생성한다.

예시의 guard는 입력 길이가 최소 2 이상이어야 한다는 조건을 만들고, matmul 예시에서는 두 차원의 일치와 함께 허용 범위를 제한한다. 즉, 컴파일된 코드는 무작정 모든 shape를 받는 것이 아니라, 실제로 안전하다고 증명된 범위 안에서만 재사용된다.

동적 shape 컴파일은 여전히 연구 성격이 강하다. 실무에서는 다음 같은 전략이 중요하다.

• padding으로 특정 bucket size에 맞추기
• vectorized fast path와 slow path를 함께 두어 leftover를 처리하기
• 자주 나오는 shape를 profile-guided optimization처럼 우선 최적화하기
• 드문 shape는 eager mode로 돌려 컴파일 비용을 아끼기

결국 핵심은, PyTorch가 shape 변화 자체를 없애려는 것이 아니라 그 변화를 제약, 재컴파일, guard로 관리해 성능과 유연성을 함께 확보한다는 점이다. 동적 shape는 까다롭지만, 적절한 compiler 지원이 있으면 대규모 배포 환경에서도 충분히 실용적으로 다룰 수 있다.