PyTorch 2.14에서 동적 shape 선언을 @dynamic_spec 하나로 묶는 방법 | DAKER 커뮤니티
PyTorch에서 동적 shape를 다룰 때 가장 번거로운 지점은 같은 모델이라도 진입점마다 선언 방식이 달랐다는 점입니다. torch.compile, torch.export, make_fx를 함께 쓰는 흐름에서는 특히 이 차이가 실험 재현성과 디버깅 난도를 높이곤 했습니다.
PyTorch 2.14는 이 문제를 줄이기 위해 @dynamic_spec과 ShapesSpec을 도입했습니다. 공식 2.14 릴리스 블로그의 Compilation and Export 절을 기준으로, 무엇이 달라졌는지와 함께 주의할 점을 간단히 정리합니다. 다만 이 API는 실험적이며 블로그에도 API Unstable로 표시되어 있습니다.
가변 차원을 함수나 모듈에 한 번 선언하면 torch.compile, torch.export, make_fx가 같은 스펙을 읽습니다.
왜 이전 방식이 헷갈렸는지
기존에는 동적 shape를 다루는 방식이 진입점마다 달랐습니다. torch.export는 보통 호출부의 dynamic_shapes 딕셔너리를 사용했고, torch.compile은 상대적으로 거친 dynamic= 플래그에 가까운 방식으로 동작했습니다. make_fx는 전역 트레이싱 모드에 의존하는 경우가 많았습니다.
이 때문에 같은 모델을 두고도 선언 위치가 제각각이었고, 배치 크기만 바뀌는 실험에서도 재컴파일이나 가드 실패가 섞여 나타나기 쉬웠습니다.
PyTorch 2.14에서 바뀐 점
2.14에서는 torch.fx.experimental.dynamic_spec 아래의 ShapesSpec으로 차원 이름, 범위, 파생 차원, 가정을 선언할 수 있습니다. 예를 들어 ShapeVar("batch", min=2, max=128)처럼 범위를 두거나, batch % 2 == 0 같은 조건을 함께 둘 수 있습니다.
이 스펙을 @dynamic_spec으로 함수나 모듈의 forward에 붙이면, 호출부에서 같은 스펙을 다시 넘기지 않아도 됩니다. 선언된 차원은 unbacked symbol로 취급되기 때문에, 트레이스 시점에 관찰한 배치 크기로 조용히 특수화되지 않는다는 점도 핵심입니다.
동적 shape 선언을 호출부가 아니라 함수나 모듈 정의 쪽으로 옮길 수 있다는 점이 2.14 변화의 핵심입니다.
성능과 해결 과정에서 살펴볼 점
동적 shape를 선언했다고 해서 모든 shape 관련 문제가 같은 방식으로 드러나는 것은 아닙니다. shape에 의존하는 분기는 가드 재컴파일이 아니라 data-dependent 오류로 나타날 수 있습니다. 이런 경우에는 먼저 분기 조건이 텐서 shape에 걸려 있지 않은지 확인하는 것이 좋습니다.
make_fx는 블로그 기준으로 tracing_mode="fake"에서만 지원이 제한적입니다. 또한 데코레이터와 호출부의 dynamic_shapes=를 함께 쓰거나, 스펙과 prefer_deferred_runtime_asserts_over_guards=True를 섞으면 오류가 납니다.
함께 보면 좋은 DAKER 학습 자료
기본기를 다시 점검하고 싶다면 텐서·디바이스·기본 연산을, 배포와 성능 관점까지 이어서 보고 싶다면 배포 감각·성능 체크를 함께 보면 맥락을 잡는 데 도움이 됩니다.
참고 자료
실제로는 어떤 지점에서 동적 shape 선언 위치가 가장 자주 헷갈렸는지 경험이 있다면 함께 나눠주셔도 좋겠습니다.