[Easy] Small DDP fixes (#72455)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/72455
- Improve helper function
- Improve/fix some logging
ghstack-source-id: 148840678
Test Plan: CI
Reviewed By: zhaojuanmao
Differential Revision: D34044865
fbshipit-source-id: d2ae820effaaaecdd7155ffa8d3a1d8ebbd9f39e
(cherry picked from commit 3efbea8f41463fc815b80f3d07ba81dd6c894e11)
