⚡️ Add torch.compile to PatchPredictor

[Abdol]

This mini-PR adds torch.compile functionality to PatchPredictor.

[measty]

At the moment the model is compiled before it is sent to the GPU (if GPU is being used). I think at least some of what torch.compile does is device-aware, so it may be better to compile after it is sent to device. Have you tried testing if the ordering makes a difference?

[shaneahmed]

At the moment the model is compiled before it is sent to the GPU (if GPU is being used). I think at least some of what torch.compile does is device-aware, so it may be better to compile after it is sent to device. Have you tried testing if the ordering makes a difference?

At the moment the model is compiled before it is sent to the GPU (if GPU is being used). I think at least some of what torch.compile does is device-aware, so it may be better to compile after it is sent to device. Have you tried testing if the ordering makes a difference?

Agree, this should be checked.

[shaneahmed]

Please enable tests for this branch.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.89%. Comparing base (b2f57ee) to head (150678b).

Additional details and impacted files

@@                  Coverage Diff                   @@
##           enhance-torch-compile     #776   +/-   ##
======================================================
  Coverage                  99.89%   99.89%           
======================================================
  Files                         69       69           
  Lines                       8578     8589   +11     
  Branches                    1641     1642    +1     
======================================================
+ Hits                        8569     8580   +11     
  Misses                         1        1           
  Partials                       8        8           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[Abdol]

At the moment the model is compiled before it is sent to the GPU (if GPU is being used). I think at least some of what torch.compile does is device-aware, so it may be better to compile after it is sent to device. Have you tried testing if the ordering makes a difference?

@measty I believe the model is compiled only when the forward function is called. See link 1 and link 2.

[Abdol]

@shaneahmed torch.compile is not compatible with Python 3.12 (see here). This has triggered an error when running CI with Python 3.12:

pytorch/pytorch#120233

if sys.version_info >= (3, 12):
raise RuntimeError("Dynamo is not supported on Python 3.12+")
E RuntimeError: Dynamo is not supported on Python 3.12+

Should we disable torch.compile for this version in the PR?

[shaneahmed]

I do not think we need this variable. I think we should only call this function if not disable

[Abdol]

Thank you @shaneahmed. That could be done, too. However, I'm mirroring the PyTorch implementation, which includes a disable flag in the function (torch.compile).

[shaneahmed]

Do we need rcparam for this? We can just set this as kwargs argument in the engines.

[Abdol]

Thank you @shaneahmed. Having kwargs for torch_compile_mode would work, too. I may suggest to keep rcParam for now until we implement it in the new engine design. Happy to discuss it in our next meeting.