Enhance fakepg: add fsdp+tp tests (#107626)
from working on a starter task with @wanchaol (T161350434):
Add the previously unsupported fsdp+tp example in FakePG, which required scatter and broadcast, as a unit test: https://github.com/pytorch/examples/blob/main/distributed/tensor_parallelism/two_d_parallel_example.py
Pull Request resolved: https://github.com/pytorch/pytorch/pull/107626
Approved by: https://github.com/wanchaol
ghstack dependencies: #107625