[Feature Request] torch compile / integrating intel extension for pytorch

[george-adams1]

🚀 Feature

Request to integrate the Intel extension for PyTorch into sb3. The Intel extension for PyTorch optimizes the PyTorch library to better utilize the computational capabilities of Intel processors. By integrating this extension, SB3 users that utilize Intel processors could potentially experience significant performance improvements.

Motivation

Maximizing the performance of Intel processors. This can lead to faster training times and more efficient resource utilization, improving the overall experience for users of SB3 who use Intel processors.

Pitch

The integration of the Intel extension for PyTorch into SB3 would involve modifying the library to utilize the extension when it detects that it is running on an Intel processor.

Alternatives

An alternative to this would be to provide guidelines on how users can manually integrate the Intel extension into their SB3 setups.

Additional context

No response

Checklist

• I have checked that there is no similar issue in the repo

[araffin]

Hello,
could you give some pointer/example on how to integrate that extension?

will you be willing to contribute such extension?

[duburcqa]

I believe it was a reference to this project: https://github.com/intel/intel-extension-for-pytorch

The latest version is compatible with torch.compile. This means that it is now much easier to leverage their optimizations without a complete rework of the libraries relying on pytorch. Here is the example script. Yet, it is describe here as an inference only backend.

[araffin]

Thanks for the pointers @duburcqa =)!

I see. In that case, it is already supported as SB3 exposes its PyTorch policy via the .policy attribute (this is a nn.Module, see doc).

Related: #1391 and #1439

You should also know that the bottleneck at train time is the gradient update, so it won't help you much to optimize inference time (unless you are at test time), we do have an experimental Jax version (SBX) if you need significant speed up (see README).

[araffin]

I would actually welcome a PR that shows how to use th.compile in our doc =)

[araffin]

As a follow up, it would be nice to investigate where to have torch compile for some speed boost (see https://github.com/pytorch-labs/LeanRL, cc @vmoens).
Given SBX experience (https://github.com/araffin/sbx), biggest speed boost should come from compiling the backward pass and combining multiple gradient steps to avoid for loops and minimizing transfer between CPU/GPU (to see the speed boost, maybe starting with CrossQ from SB3 contrib).

[vmoens]

Hey
Happy to help with this.
Compiling several steps together isn't a bad idea.
The only thing that usually works best is this:

• put all the updates in a single update function, including zero_grad, forward, backward, clip grad norm, step and target params updates
• if possible group optimizers together (this can be done using param_groups in they all have the same number of steps)
• have a look at the graph breaks using TORCH_LOGS="+graph_breaks,recompiles" python code.py
• avoid cpu-gpu transfers in the compiled region, and use torch whenever possible (not numpy).
• prefer torch.where to masking, avoid control flow if can be avoided (again sometimes torch.where can help)
• use mode="reduce-overhead" or tensordict's CudaGraphModule like we do in leanrl

That's the 101 but I'd be happy to review or help implement a compiled version of any script.