[Reland] Add tests for DDP control flow models. (#47470)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/47470
Reland of https://github.com/pytorch/pytorch/pull/47206, which was reverted due to failing multigpu tests.
The fix to make multigpu tests work is to compare against `torch.tensor([world_size, 0])`, not hardcode `torch.tensor([2, 0]` which assumes a world size of 2.
Original commit description:
As discussed offline with pritamdamania87, add testing to ensure per-iteration and rank-dependent control flow works as expected in DDP with find_unused_parameters=True.
ghstack-source-id: 115993934
ghstack-source-id: 115993934
Test Plan: CI
Reviewed By: pritamdamania87
Differential Revision: D24767893
fbshipit-source-id: 7d7a2449270eb3e72b5061694e897166e16f9bbc
