[shard] Extensible Sharder and ShardingPlanner (#75844)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/75844
Add extensible Sharder and ShardingPlanner for advanced users to extend and do custom sharding, some details:
* ShardingPlanner is merely a interface that could be extended to plan advanced sharding strategies, it will produce a ShardingPlan that outlines the detailed sharding strategies for each submodule, note that it could also include custom sharders.
* a custom Sharder is another interface that could be extended to define advanced sharding strategies for specific modules (i.e. a collection of `nn.EmbeddingBag`s.) The custom sharder defined by the user can be picked up by `share_module` API. `shard_module` will call `Sharder.shard` to the module, the output of the `Sharder.shard` will replace the original submodule.
ghstack-source-id: 154978919
Test Plan:
test_sharder
test_sharding_plan
Reviewed By: pritamdamania87
Differential Revision: D35658657
fbshipit-source-id: ad9deb39a6a0ab91ceb1de3a78c4709aa50ad699
(cherry picked from commit f9796471bd7f8ad94a9d3462cac30eacacf68b31)
