Support regional compilation

[chaojun-zhang]

What does this PR do?

DeepSpeed's addition of forward and backward hooks to the model causes torch dynamo to fall back to eager mode during compilation for the backward pass (e.g., using CppTensorPreHookFunction), leading to a significant performance degradation.

By applying regional compilation to the model, we can effectively isolate DeepSpeed's hooks from the backward graph, allowing dynamo to successfully trace the backward pass of the model, thus improving performance.

More info about regional compilation:

https://pytorch.org/tutorials/recipes/regional_compilation.html
https://pytorch.org/blog/pytorch2-5/

Fixes # (issue)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

[chaojun-zhang]

@regisss can you have a look ?