[DDP] Add a test case to test a larger model (#4085)
Summary:
This commit adds a test case to test a larger model that can trigger
multiple all_reduces instead of one. It also fixes a xmp issue while
the launcher being used to run consecutive mp experiments.
Test Plan:
XRT:
MASTER_ADDR=localhost MASTER_PORT=6000 python test/test_ddp.py TestXrtDistributedDataParallel.test_ddp_correctness_large_net
PJRT:
PJRT_DEVICE=TPU python test/pjrt/test_ddp.py TestPjRtDistributedDataParallel.test_ddp_correctness_large_net
