Repository: pytorch
Issue: Umbrella issue for Dynamo tracing through ProcessGroup related torch.distributed APIs
Labels: oncall: distributed,triaged,oncall: pt2,pt2d-triage-nov2024

### 🐛 Describe the bug

Internal xref: https://www.internalfb.com/diff/D56124045?dst_version_fbid=1206743993978310

It looks like a lot of basic APIs in torch.distributed are not Dynamo traceable, which we noticed when we were trying to add some logic for checking up on PGs in torchrec. For the most part, I'm expecting these to be constant propagated.

torch.distributed.get_backend does not work:

```
Failures:
  1) torchrec.distributed.tests.test_pt2_multiprocess.TestPt2Train: test_compile_multiprocess_fake_pg
    1) Unsupported: builtin: eq [<class 'torch._dynamo.variables.distributed.ProcessGroupVariable'>, <class 'torch._dynamo.variables.misc.GetAttrVariable'>] False
    
    from user code:
       File "/data/users/hhy/fbsource/buck-out/v2/gen/fbcode/a009654a71cfe201/torchrec/distributed/tests/__test_pt2_multiprocess__/test_pt2_multiprocess#link-tree/torchrec/distributed/model_parallel.py", line 295, in forward
        return self._dmp_wrapped_module(*args, **kwargs)
      File "/data/users/hhy/fbsource/buck-out/v2/gen/fbcode/a009654a71cfe201/torchrec/distributed/tests/__test_pt2_multiprocess__/test_pt2_multiprocess#link-tree/torch/nn/modules/module.py", line 1747, in _call_impl
        return forward_call(*args, **kwargs)
      File "/data/users/hhy/fbsource/buck-out/v2/gen/fbcode/a009654a71cfe201/torchrec/distributed/tests/__test_pt2_multiprocess__/test_pt2_multiprocess#link-tree/torchrec/distributed/tests/test_pt2_multiprocess.py", line 165, in forward
        kt: KeyedTensor = self._ebc(x)
      File "/data/users/hhy/fbsource/buck-out/v2/gen/fbcode/a009654a71cfe201/torchrec/distributed/tests/__test_pt2_multiprocess__/test_pt2_multiprocess#link-tree/torch/nn/modules/module.py", line 1747, in _call_impl
        return forward_call(*args, **kwargs)
      File "/data/users/hhy/fbsource/buck-out/v2/gen/fbcode/a009654a71cfe201/torchrec/distributed/tests/__test_pt2_multiprocess__/test_pt2_multiprocess#link-tree/torchrec/distributed/types.py", line 881, in forward
        dist_input = self.input_dist(ctx, *input, **kwargs).wait().wait()
      File "/data/users/hhy/fbsource/buck-out/v2/gen/fbcode/a009654a71cfe201/torchrec/distributed/tests/__test_pt2_multiprocess__/test_pt2_multiprocess#link-tree/torchrec/distributed/embeddingbag.py", line 1090, in input_dist
        awaitables.append(input_dist(features_by_shard))
      File "/data/users/hhy/fbsource/buck-out/v2/gen/fbcode/a009654a71cfe201/torchrec/distributed/tests/__test_pt2_multiprocess__/test_pt2_multiprocess#link-tree/torch/nn/modules/module.py", line 1747, in _call_impl
        return forward_call(*args, **kwargs)
      File "/data/users/hhy/fbsource/buck-out/v2/gen/fbcode/a009654a71cfe201/torchrec/distributed/tests/__test_pt2_multiprocess__/test_pt2_multiprocess#link-tree/torchrec/distributed/sharding/tw_sharding.py", line 259, in forward
        return self._dist(sparse_features)
      File "/data/users/hhy/fbsource/buck-out/v2/gen/fbcode/a009654a71cfe201/torchrec/distributed/tests/__test_pt2_multiprocess__/test_pt2_multiprocess#link-tree/torch/nn/modules/module.py", line 1747, in _call_impl
        return forward_call(*args, **kwargs)
      File "/data/users/hhy/fbsource/buck-out/v2/gen/fbcode/a009654a71cfe201/torchrec/distributed/tests/__test_pt2_multiprocess__/test_pt2_multiprocess#link-tree/torchrec/distributed/dist_data.py", line 637, in forward
        return KJTAllToAllSplitsAwaitable(
      File "/data/users/hhy/fbsource/buck-out/v2/gen/fbcode/a009654a71cfe201/torchrec/distributed/tests/__test_pt2_multiprocess__/test_pt2_multiprocess#link-tree/torchrec/distributed/dist_data.py", line 495, in __init__
        self._splits_awaitable = SplitsAllToAllAwaitable(
      File "/data/users/hhy/fbsource/buck-out/v2/gen/fbcode/a009654a71cfe201/torchrec/distributed/tests/__test_pt2_multiprocess__/test_pt2_multiprocess#link-tree/torchrec/distributed/dist_data.py", line 242, in __init__
        if dist.get_backend(pg) == "fake":
      File "/data/users/hhy/fbsource/buck-out/v2/gen/fbcode/a009654a71cfe201/torchrec/distributed/tests/__test_pt2_multiprocess__/test_pt2_multiprocess#link-tree/torch/distributed/distributed_c10d.py", line 1214, in get_backend
        if _rank_not_in_group(pg):
      File "/data/users/hhy/fbsource/buck-out/v2/gen/fbcode/a009654a71cfe201/torchrec/distributed/tests/__test_pt2_multiprocess__/test_pt2_multiprocess#link-tree/torch/distributed/distributed_c10d.py", line 895, in _rank_not_in_group
        return group == GroupMember.NON_GROUP_MEMBER
```

But it would also be useful to hose down the rest of the public torch.distributed API and make sure they work (when it would be trivial to make them work)

### Versions

main

cc @XilunWu @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @chauhang @penguinwu