Repository: pytorch
Issue: [PT2] AOTAutograd de-dups but skips de-dup guards for DDP
Labels: oncall: distributed,triaged,module: ddp,oncall: pt2,pt2d-triage-nov2024

Run DDP with a shared buffer (different TorchDynamo `Source`):
<details>
<summary> Repro Script </summary>

```
"""
torchrun --standalone --nproc_per_node=1 test/dup_repro.py
TORCH_LOGS=aot,dynamo torchrun --standalone --nproc_per_node=1 test/dup_repro.py
"""
import os

import torch
import torch.distributed as dist
import torch.nn as nn
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.nn.parallel import DistributedDataParallel as DDP

USE_FSDP = False
print(f"USE_FSDP={USE_FSDP}")

os.environ["TORCHDYNAMO_PRINT_GUARDS"] = "1"


class BufModule(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.register_buffer(
            "_buf", torch.randn((3,), requires_grad=False, device="cuda")
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x + self._buf


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        # Define a parameter since DDP requires at least one
        self._param = nn.Parameter(torch.randn((1,), device="cuda"))
        self._buf_module = BufModule()
        # Use same tensor but with different source
        self.register_buffer("_buf", self._buf_module._buf)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        z = x + self._buf
        z = self._buf_module(z)
        z += self._param
        return z


dist.init_process_group(backend="nccl")
gpu_id = int(os.environ["LOCAL_RANK"])
device = f"cuda:{gpu_id}"
torch.cuda.set_device(device)

model = Model()
if USE_FSDP:
    model = FSDP(model, use_orig_params=True)
else:
    model = DDP(model, device_ids=[dist.get_rank()])
model = torch.compile(model)
if USE_FSDP:
    assert model._buf is model._buf_module._buf
else:
    assert model.module._buf is model.module._buf_module._buf
inp = torch.randn((2, 3), device="cuda")
model(inp)
```
</details>

DDP forward graph:
```
====== Forward graph 0 ======
 <eval_with_key>.7 class GraphModule(torch.nn.Module):
    def forward(self, primals_1: f32[1], primals_2: f32[3], primals_3: f32[2, 3]):
        # File: test/dup_repro.py:39, code: z = x + self._buf
        add: f32[2, 3] = torch.ops.aten.add.Tensor(primals_3, primals_2);  primals_3 = None
        
        # File: test/dup_repro.py:27, code: return x + self._buf
        add_1: f32[2, 3] = torch.ops.aten.add.Tensor(add, primals_2);  add = primals_2 = None
        
        # File: test/dup_repro.py:41, code: z += self._param
        add_2: f32[2, 3] = torch.ops.aten.add.Tensor(add_1, primals_1);  add_1 = primals_1 = None
        return [add_2]
```
It looks like `model._buf` and `model._buf_module._buf` got de-duplicated since the forward graph only has `primals_2` as the size `[3]` tensor (`primals_1` is `model._param` and `primals_3` is the input tensor).

However, there is no de-dup guard:
```
GUARDS ___guarded_code.valid
and ___check_obj_id(L['self'], 140390721572928)
and L['self'].training == True
and not ___are_deterministic_algorithms_enabled()
and ___check_type_id(G['__import_torch_dot_nn_dot_modules_dot_module']._global_forward_hooks, 93969908371680)
and set(G['__import_torch_dot_nn_dot_modules_dot_module']._global_forward_hooks.keys()) == set()
and ___check_type_id(G['__import_torch_dot_nn_dot_modules_dot_module']._global_backward_hooks, 93969908371680)
and set(G['__import_torch_dot_nn_dot_modules_dot_module']._global_backward_hooks.keys()) == set()
and ___check_type_id(G['__import_torch_dot_nn_dot_modules_dot_module']._global_forward_pre_hooks, 93969908371680) 
and set(G['__import_torch_dot_nn_dot_modules_dot_module']._global_forward_pre_hooks.keys()) == set()
and ___check_type_id(G['__import_torch_dot_nn_dot_modules_dot_module']._global_backward_pre_hooks, 93969908371680) 
and set(G['__import_torch_dot_nn_dot_modules_dot_module']._global_backward_pre_hooks.keys()) == set()
and ___check_tensors(L['x'])
```
I was expecting a guard like:
```
L['self']._buf is L['self']._buf_module._buf
```

Let me know if I am not understanding correctly.

cc @H-Huang @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @chauhang @penguinwu @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @XilunWu @tianyu-l @yf225 @ezyang @msaroufim @bdhirsh @anijain2305 @zou3519 @soumith @ngimel