torch.compile의 Tracing과 Compile 시점에 대해 문의드립니다.

[osehyeon]

안녕하세요, PyTorch + NPU LAB에 참가 중인 오세현입니다.

torch.compile을 사용하려면 모델을 래핑하여 compiled_model을 생성하고, 이를 기존 모델과 동일한 방식으로 사용할 수 있는 것으로 알고 있습니다.

compiled_model = torch.compile(model) # [1]
outputs = compiled_model(images) # [2]

여기서 다음 사항이 궁금합니다:

1. Tracing과 Compile은 언제 이루어지나요?
   Tracing과 Compile이 [1]의 torch.compile() 호출 시점에서 Fake Tensor를 사용하여 이루어지는 것인지, 아니면 [2]의 모델이 첫 번째 입력을 처리하는 시점에서 이루어지는 것인지, 혹은 다른 방식으로 이루어지는 것인지 궁금합니다.

감사합니다.

[appleparan]

저도 같이 공부하는 입장에서 생각해보자면,

• torch.jit.trace는 data-dependent flow를 제대로 처리하지 못하는 것으로 알고있습니다. 실질적으로 데이터가 들어갔을 때 결정된다고 생각되므로 [2]에서 처리가 되지 않을까요?
• torch.jit.script는 이런 data dependent한 flow는 처리하는 대신 numpy같은 non-PyTorch 함수를 제대로 처리하지 못한다고 알고있습니다.
• 이런 Torch Scripting과 FX Tracing은 이 문서를 참고하면 좋을 것 같습니다.
• torch.compile은 PEP 523덕분에 TorchDynamo를 이용하여 non-PyTorch 함수를 잘 처리하는 것으로 알고있습니다. TorchDynamo를 봤을 때, JIT으로 이루어지기 때문에 실질적으로 데이터가 들어왔을 때 bytecode analysis를 통해 컴파일이 이루어지지 않나 생각됩니다.