Skip to content

Add optional torch.compile support to InferenceRunner via InferenceConfig - #3098

Closed
maximpavliv wants to merge 6 commits into
mainfrom
maxim/torch_compile
Closed

Add optional torch.compile support to InferenceRunner via InferenceConfig#3098
maximpavliv wants to merge 6 commits into
mainfrom
maxim/torch_compile

Conversation

@maximpavliv

@maximpavliv maximpavliv commented Sep 17, 2025

Copy link
Copy Markdown
Contributor

Description

This PR adds optional support for torch.compile in the InferenceRunner class via the new InferenceConfig class (built on top of PR #3104).

If compilation is enabled, the model is compiled at InferenceRunner initialization using the specified backend. In our speed benchmarking suite, enabling torch.compile shows faster inference speed:
concat

However, there are some important caveats:

  • Compilation may fail depending on the user’s environment or setup.
  • It introduces a one-time compilation overhead during inference initialization, which may be confusing.
  • Our current CI/CD tests fail when compilation is enabled (further investigation needed).

For these reasons, torch.compile is disabled by default, but advanced users can opt in at their own risk.

Key changes

  • Added torch.compile support via the InferenceConfig class.

  • Compilation parameters can now be passed either through the InferenceConfig object or via the inference: block in pytorch_config.yaml, consistent with the multithreading configuration.

  • When enabled, the model is compiled at initialization using the specified backend.

Notes

  • Default behavior remains unchanged for all users.
  • Advanced users can enable compilation for potential inference speedups.
  • This PR builds on the multithreading and InferenceConfig infrastructure introduced in PR #3104.

@maximpavliv maximpavliv added the WORK IN PROGRESS! developers are currently working on this feature... stay tuned. label Sep 17, 2025
@maximpavliv
maximpavliv marked this pull request as draft September 17, 2025 13:33
@maximpavliv
maximpavliv force-pushed the maxim/torch_compile branch 2 times, most recently from 8e1469a to c732fe7 Compare September 17, 2025 19:34
@maximpavliv maximpavliv changed the title [WIP] Add torch.compile to InferenceRunner Add optional torch.compile support and inference configuration classes Sep 18, 2025
@maximpavliv maximpavliv added new feature DLC3.0🔥 speed-checking and removed WORK IN PROGRESS! developers are currently working on this feature... stay tuned. labels Sep 18, 2025
@maximpavliv
maximpavliv marked this pull request as ready for review September 18, 2025 12:29
@maximpavliv maximpavliv changed the title Add optional torch.compile support and inference configuration classes Add optional torch.compile support and inference configuration classes Sep 18, 2025
@maximpavliv
maximpavliv marked this pull request as draft September 23, 2025 08:43
@maximpavliv maximpavliv changed the title Add optional torch.compile support and inference configuration classes Add torch.compile support via CompileConfig in InferenceConfig for optional model compilation Sep 23, 2025
@maximpavliv maximpavliv changed the title Add torch.compile support via CompileConfig in InferenceConfig for optional model compilation Add optional torch.compile support to InferenceRunner via InferenceConfig Sep 23, 2025
@maximpavliv
maximpavliv marked this pull request as ready for review September 24, 2025 08:15
@arashsm79

Copy link
Copy Markdown
Contributor

Tested this on various projects. After the initial compilation delay, the speed up is great!
Some notes for future:

  • Would be nice to use the dot notation: {"compile.enabled": False}
  • Would be better to error out if the user has input the wrong key. For example:
    inference_cfg={"compile": {"enble": False}} should fail.
  • Torch compile kept failing for CTD models. Maybe we should add a note for this.

@maximpavliv

Copy link
Copy Markdown
Contributor Author

Thanks for your review @arashsm79 !

As this PR no longer has any diffs against main because its commits were included in #3105, I'm closing as redundant.

I'll address your feedback in #3106 @arashsm79 .

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants