Enable passing mode and other kwargs to torch.compile in BaseHandler

[eballesteros]

🚀 The feature

Enable passing mode and other kwargs to torch.compile in BaseHandler

At the moment only backend is configurable.

Motivation, pitch

I'm working on a custom handler that re-uses initialize from BaseHandler. It would be great to have more control over torch.compile options.

Alternatives

I've considered re-implementing initialize in my custom handler, but it seems a bit wasteful just for this.

Additional context

Happy to take a stab at it myself

[msaroufim]

This makes sense, if you'd like to take a stab at it I'd be happy to review. You can take a look at how backend is implemented and do the same for mode

Ideally we'd also want to pass in an arbitrary kwarg dict config so if you're interested in that too lmk

[eballesteros]

Sounds good! Let's go for the arbitrary kwarg dict since we're at it. After taking a look at how backend is implemented I think it should be pretty straight forward, but wanted to run the API by you before jumping to code:

I was thinking about optionally adding a second key to model_config.yaml named compile_options, or something along those lines, with a dict value. Something like:

pt2: "inductor"
compile_options: {"mode": "reduce-overhead", "fullgraph": True}

The other option I can think of without breaking backwards compatibility is allowing the pt2 key to have either string or dict value. If it's string we treat that as the backend just like now, if it's dict we pass it as kwargs.

Let me know if either of those sound good, or if you had a better implementation in mind.

[msaroufim]

The other option I can think of without breaking backwards compatibility is allowing the pt2 key to have either string or dict value. If it's string we treat that as the backend just like now, if it's dict we pass it as kwargs.

This sounds best to me