[PERF] Adds DDP mode to imagenet and mnist tests (#4135)
Summary:
This pull requests enhance test_train_mp_imagenet.py and test_train_mp_mnist.py
with the ability to use torch.nn.parallel.DistributedDataParallel. So it's
easier for us to benchmark it and compare with our native approach.
Test Plan:
MASTER_ADDR=localhost MASTER_PORT=6000 python test/test_train_mp_imagenet.py --ddp --fake_data
MASTER_ADDR=localhost MASTER_PORT=6000 python test/test_train_mp_mnist.py --ddp --fake_data
