pytorch

f5d231d5 - move rebuild buckets from end of first iteration to beginning of second iteration (#44326)

move rebuild buckets from end of first iteration to beginning of second iteration (#44326) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/44326 Part of relanding PR #41954, this refactoring is to move rebuild_buckets call from end of first iteration to beginning of second iteration ghstack-source-id: 112011490 Test Plan: unit tests Reviewed By: mrshenli Differential Revision: D23583017 fbshipit-source-id: ef67f79437a820d9b5699b651803622418499a83