We need a custom `save` function for checkpointing a lazy model, similar to what exists in PyTorch/XLA:
3eb8a9d9eb/torch_xla/core/xla_model.py (L994)
The purpose of this function is to move any lazy tensors to CPU before saving the checkpoint.
The way I implemented it was to create a general structure visitor, adapted from a function that we use quite often in Cerebras internal repositories. If there is a better tool already available in PyTorch that does the same things, I'm open to suggestions.
CC: @wconstab @Krovatkin @JackCaoG
Pull Request resolved: https://github.com/pytorch/pytorch/pull/83294
Approved by: https://github.com/wconstab
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/75046
Merge the code needed for dynamic+ltc integration from the staging branch to the master branch.
Test Plan:
Unit test
```
pytest test_extract_compiled_graph
```
test thru dynamo
```
LTC_TS_CUDA=1 time python torchbench.py --speedup-ltc -dcuda --nvfuser --randomize-input --only <model name>
```
Reviewed By: alanwaketan
Differential Revision: D35300646
Pulled By: shunting314
fbshipit-source-id: 09ed20d3bb8ef80e4b93ba87ea3356a07d2dccdb
(cherry picked from commit 2b56771cdfd2cfa825c65ee9fd42338fb372fb32)
Summary:
This adds a minimal set of python bindings for lazy tensor and the torchscript backend.
It targets the APIs that are used by the `test_ts_opinfo.py` test (which it also lands, from lazy_tensor_staging, where it is [lazy_tensor_core/test/test_lazy.py](https://github.com/pytorch/pytorch/blob/lazy_tensor_staging/lazy_tensor_core/test/test_lazy.py)).
We should land more python bindings obviously. I just wanted to focus on a minimal set that can also be tested, and use it to agree on how we organize the bindings, then others could easily contribute bindings on top of this infrastructure.
cc JackCaoG
Pull Request resolved: https://github.com/pytorch/pytorch/pull/74508
Reviewed By: pbelevich
Differential Revision: D35032152
Pulled By: wconstab
fbshipit-source-id: 526505ab355b7ad27037ece0ff814b2a4b69f1e2
(cherry picked from commit b4f73dd147472cb38003204aff228087c0230fda)
