DTensor: add comm tests to test_tp_examples (#121669)
This adds some basic comm tests to test_tp_examples. This validates that the expected distributed calls are being made for `test_transformer_training`.
Fixes #121649
Test plan:
```
pytest test/distributed/tensor/parallel/test_tp_examples.py -k test_transformer_training
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/121669
Approved by: https://github.com/wanchaol
