add sparse_nn_partition (#46390)
Summary:
WIP: This PR adds sparse_nn_partition into Partitioner class. It includes logical device assignment for all dag nodes. The basic idea is to do size_based_partition separately for embedding nodes and non-embedding nodes. A test unit is also added in test_fx_experimental.py.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/46390
Reviewed By: gcatron
Differential Revision: D24555415
Pulled By: scottxu0730
fbshipit-source-id: 8772af946d5226883759a02a1c827cfdfce66097
