[Dist Profiling] Support shape recording for profiling collectives (#51822)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/51822
Adds support for shape recording for profiling distributed collectives, for nccl/gloo backends. Added
both cpp and python tests to ensure that shapes are recorded properly. Note that we don't add `ProcessGroupNCCLTest`s since they need to be modified to support single process per device and > 1 world size.
ghstack-source-id: 121507509
Test Plan: CI
Reviewed By: mrzzd
Differential Revision: D26291739
fbshipit-source-id: 5f7bd54d8c36d17a4a29e172b25266ca3dbd8fbd
