Enable creation and transfer of ScriptModule over RPC (#48293)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/48293
This PR enables a `ScriptModule` to be created on a remote worker, retrieve it to current device, and run methods on it remotely.
In order to do this, we define custom pickling for `ScriptModule` in our `InternalRPCPickler`. The pickling basically uses torch.save and torch.load to save/recover the ScriptModule.
We test that we can create and retrieve a ScriptModule with rpc_sync, and create RRef to ScriptModule with rpc.remote. We can also run remote methods on the rref and transfer it to current worker.
Although we can run methods remotely on the RRef to ScriptModule, this does not currently work with RRef helper, filed to track that.
ghstack-source-id: 117275954
Test Plan: CI
Reviewed By: wanchaol
Differential Revision: D25107773
fbshipit-source-id: daadccf7bd25fe576110ee6e0dba6ed2bcd3e7f3
