feat: torch 2.0 #3682
feat: torch 2.0 #3682
Conversation
Codecov Report
@@ Coverage Diff @@
## main #3682 +/- ##
==========================================
- Coverage 31.74% 31.71% -0.04%
==========================================
Files 149 149
Lines 12149 12162 +13
Branches 2001 2003 +2
==========================================
Hits 3857 3857
- Misses 8008 8021 +13
Partials 284 284
|
Signed-off-by: Aaron <29749331+aarnphm@users.noreply.github.com>
| class ModelOptions(PartialKwargsModelOptions): | ||
| fullgraph: bool = False | ||
| dynamic: bool = False | ||
| backend: t.Union[str, t.Callable[..., t.Any]] = "inductor" | ||
| mode: t.Optional[str] = None | ||
| options: t.Optional[t.Dict[str, t.Union[str, int, bool]]] = None | ||
| disable: bool = False | ||
|
|
||
|
|
There was a problem hiding this comment.
First, let’s call this PytorchOptions and only when import to bentoml.pytorch we rename it to ModelOptions.
Second, I think maybe it’s better to have single compile_kwargs dict in PytorchOptions instead of polluting the name space of PytorchOptions. Maybe we will have only 2 entries:
enable_compilecompile_kwargs
| def load_model( | ||
| bentoml_model: str | Tag | Model, | ||
| device_id: t.Optional[str] = "cpu", | ||
| **compile_kwargs: t.Any, |
There was a problem hiding this comment.
I don’t think we need do compile at load_model level. Let’s just save the original model and do torch.compile when init the runner (if user set enable_compile=True)
| opts = t.cast(ModelOptions, bento_model.info.options) | ||
| if get_pkg_version("torch") >= "2.0.0": | ||
| _load_model = partial( | ||
| load_model, | ||
| fullgraph=opts.fullgraph, | ||
| dynamic=opts.dynamic, | ||
| backend=opts.backend, | ||
| mode=opts.mode, | ||
| options=opts.options, | ||
| disable=opts.disable, | ||
| ) | ||
| else: | ||
| _load_model = load_model |
There was a problem hiding this comment.
I think we can just model = load_model(…), and then torch.compile(model, **compile_kwargs) if enable_compile is True.
This PR aims to bring Torch 2.0 support to
save_modelviatorch.compile