Be fix shard inbalance (#60206)
Summary:
First step to address https://github.com/pytorch/pytorch/issues/60136
Pull Request resolved: https://github.com/pytorch/pytorch/pull/60206
Reviewed By: janeyx99
Differential Revision: D29215237
Pulled By: walterddr
fbshipit-source-id: ec25beb57366ef2eaf37878cdea391b245de9bef
Author
Rong Rong (AI Infra)
