[Dynamo + DDP] If DDP partitions FX graph generated by Dynamo correctly
Labels: triaged, oncall: pt2, module: dynamo

### 🐛 Describe the bug

Repro:
```
import torch
import logging
import torch._dynamo

torch._logging.set_logs(dynamo=logging.DEBUG, bytecode=True)

import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(420)

class MyModel(nn.Module):

    def __init__(self):
        super(MyModel, self).__init__()
        self.linear = torch.nn.Linear(5, 5)

    def forward(self, x):
        x = x + 1
        with torch.cuda.amp.autocast(dtype=torch.float16):
            x = self.linear(x)
            x = torch.sin(x)
        x = torch.cos(x)
        x = x - 1
        return x

x = torch.randn(5, 5, device="cuda")

m = MyModel().to("cuda")
print(m(x))

opt_m = torch.compile(backend="eager")(m)
print(opt_m(x))
```

Generated FX graph:
```
 ===== __compiled_fn_0 =====
 <eval_with_key>.0 class GraphModule(torch.nn.Module):
    def forward(self, L_x_ : torch.Tensor):
        l_x_ = L_x_
        
        # File: /scratch/ybliang/work/repos/debug/debug6.py:19, code: x = x + 1
        add = l_x_ + 1;  l_x_ = None
        
        # No stacktrace found for following nodes
        enter_functional_autocast = torch__dynamo_variables_ctx_manager_enter_functional_autocast('cuda', torch.float16, True, True)
        
        # File: /scratch/ybliang/work/repos/debug/debug6.py:21, code: x = self.linear(x)
        l__self___linear = self.L__self___linear(add);  add = None
        
        # File: /scratch/ybliang/work/repos/debug/debug6.py:22, code: x = torch.sin(x)
        sin = torch.sin(l__self___linear);  l__self___linear = None
        
        # No stacktrace found for following nodes
        exit_functional_autocast = torch__dynamo_variables_ctx_manager_exit_functional_autocast(enter_functional_autocast);  enter_functional_autocast = None
        
        # File: /scratch/ybliang/work/repos/debug/debug6.py:23, code: x = torch.cos(x)
        cos = torch.cos(sin);  sin = None
        
        # File: /scratch/ybliang/work/repos/debug/debug6.py:24, code: x = x - 1
        sub = cos - 1;  cos = None
        return (sub,)
```

Generated compiled function:
```
 __compiled_fn_0 <eval_with_key>.0 opcode         name                       target                                                  args                                 kwargs
-------------  -------------------------  ------------------------------------------------------  -----------------------------------  --------
placeholder    l_x_                       L_x_                                                    ()                                   {}
call_function  add                        <built-in function add>                                 (l_x_, 1)                            {}
call_function  enter_functional_autocast  <function enter_functional_autocast at 0x7f7594c66550>  ('cuda', torch.float16, True, True)  {}
call_module    l__self___linear           L__self___linear                                        (add,)                               {}
call_function  sin                        <built-in method sin of type object at 0x7f798f3f19e0>  (l__self___linear,)                  {}
call_function  exit_functional_autocast   <function exit_functional_autocast at 0x7f7594c74d30>   (enter_functional_autocast,)         {}
call_function  cos                        <built-in method cos of type object at 0x7f798f3f19e0>  (sin,)                               {}
call_function  sub                        <built-in function sub>                                 (cos, 1)                             {}
output         output                     output                                                  ((sub,),)                            {}
```

### Versions

N/A

cc @ezyang @gchanan @zou3519 @msaroufim @wconstab @ngimel @bdhirsh @anijain2305 @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @desertfire