[FX] Fix an issue in constructing GraphModule
Labels: fb-exported, release notes: fx

Summary:
During the enablement of IG HSTU models, we noticed an issue around model split, which the same submodule is presented in the graph modules of two split modules. (P812849675) We examined the graph of the problematic graph module, and found the graph is correct as expected (P815770164), the nodes representing the misplaced submodule is not presented in the graph, but somehow when constructing the graph module from the graph, those submodules are inserted into the module. This causing many issues when we transform the model downstream.

Further debugging shows there's some edge case issue when constructing the graph module. We created a minimal repro the issue in N4232266

Specifically, this happens when we create the submodules when constructing the graph module. It traversing the nodes in the graph, if the node's op is `call_module` or `get_attr`, the submodule from the root module is copied to the graph module.

For example, if we have simple modules like this
```
class M(torch.nn.Module):
    def __init__(self, is_train: bool = False) -> None:
        super().__init__()
        self._is_train = is_train
        self.linear1 = torch.nn.Linear(2, 2)
        self.relu1 = torch.nn.ReLU()
        self.linear2 = torch.nn.Linear(2, 2)
        self.relu2 = torch.nn.ReLU()
        self.linear3 = torch.nn.Linear(2, 2)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self._is_train:
            x = self.relu2(self.linear2(self.relu1(self.linear1(x))))
        return self.linear3(x)

class Model(torch.nn.Module):
    def __init__(self, is_train: bool = False) -> None:
        super().__init__()
        self._m = M(is_train)
        self.linear = torch.nn.Linear(2, 2)
        self.relu = torch.nn.ReLU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.relu(self.linear(x))
        x = self._m(x)
        return x

m = Model()
```
FX tracing the model `m` would get the graph
```
opcode       name        target      args           kwargs
-----------  ----------  ----------  -------------  --------
placeholder  x           x           ()             {}
call_module  linear      linear      (x,)           {}
call_module  relu        relu        (linear,)      {}
call_module  _m_linear3  _m.linear3  (relu,)        {}
output       output      output      (_m_linear3,)  {}
```
and graph module constructed as
```
Model(
  (linear): Linear(in_features=2, out_features=2, bias=True)
  (relu): ReLU()
  (_m): Module(
    (linear3): Linear(in_features=2, out_features=2, bias=True)
  )
)
```
The `linear1/2`, `relu1/2` in the `M` are not the in the graph, thus omitted in the graph module.

Now we have some use case that the module `Model`'s `forward` calling into some other functions like
```
torch.fx.wrap
def some_func(module: torch.nn.Module, x:torch.Tensor) -> torch.Tensor:
    # do something...
    return x

class Model(torch.nn.Module):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.relu(self.linear(x))
        x = self._m(x)
        x = some_func(self._m, x)
        return x
```
The calling of `some_func` add a `get_attr` node in the graph, the graph becomes
```
opcode         name        target                                  args              kwargs
-------------  ----------  --------------------------------------  ----------------  --------
placeholder    x           x                                       ()                {}
call_module    linear      linear                                  (x,)              {}
call_module    relu        relu                                    (linear,)         {}
call_module    _m_linear3  _m.linear3                              (relu,)           {}
get_attr       _m          _m                                      ()                {}
call_function  some_func   <function some_func at 0x7f7acf35e280>  (_m, _m_linear3)  {}
output         output      output                                  (some_func,)      {}
```
So when constructing the graph module, it loops through all the `get_attr` and `call_module` nodes in the graph, and insert the submodules to the graph module
- when processing node `_m.linear3`, the `_copy_attr` first checks whether `_m` is presented, inserted `_m` module, then copied `linear3` under `_m`
- then it sees the `get_attr` node `_m`, and copied the original `_m` module over to the graph module, including the submodules not in the graph

thus we get the following wrong graph module
```
Model(
  (linear): Linear(in_features=2, out_features=2, bias=True)
  (relu): ReLU()
  (_m): M(
    (linear1): Linear(in_features=2, out_features=2, bias=True)
    (relu1): ReLU()
    (linear2): Linear(in_features=2, out_features=2, bias=True)
    (relu2): ReLU()
    (linear3): Linear(in_features=2, out_features=2, bias=True)
  )
)
```

This is just the demonstration of the issue, in practice, we ended up to have the modules not in the merge graph copied to merge net in the model split, and causing various downstream issues.

To fix, we need to do some extra check when copying submodule. If the `field` already presented in the `to_module` and the node's op is not `call_module`, then we should skip copying the field.

Test Plan:
With the diff, constructing the graph module from the example
```
opcode         name        target                                  args              kwargs
-------------  ----------  --------------------------------------  ----------------  --------
placeholder    x           x                                       ()                {}
call_module    linear      linear                                  (x,)              {}
call_module    relu        relu                                    (linear,)         {}
call_module    _m_linear3  _m.linear3                              (relu,)           {}
get_attr       _m          _m                                      ()                {}
call_function  some_func   <function some_func at 0x7f0f459f0ca0>  (_m, _m_linear3)  {}
output         output      output                                  (some_func,)      {}

# graph module
Model(
  (linear): Linear(in_features=2, out_features=2, bias=True)
  (relu): ReLU()
  (_m): Module(
    (linear3): Linear(in_features=2, out_features=2, bias=True)
  )
)
```
noticed the correct graph module is constructed

try out on IG HSTU model, the merge net's graph module is constructed correctly P818433831, unblocked the enablement

Differential Revision: D48983396