Repository: pytorch
Issue: serialization of PT2E model impacts torch.fx.passes.utils.source_matcher_utils.get_source_partitions 
Labels: export-triage-review,oncall: export

### 🐛 Describe the bug

**My expectation is that get_source_partitions should act the same no matter the PT2E model is saved/loaded or not, and I have some other questions inline, thanks.**

### **case 1:**   
call get_source_partitions on the exported model without save/load.

test code:
```
import torch
from torch.fx.passes.utils.source_matcher_utils import (
                    SourcePartition,
                    get_source_partitions,
)

class M(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = torch.nn.Linear(3, 3)
        self.relu = torch.nn.ReLU()
        self.linear2 = torch.nn.Linear(3, 5)

    def forward(self, x):
        x = self.linear1(x)
        x = self.linear1(x)
        x = self.relu(x)
        x = self.linear2(x)
        return x

inputs = (torch.randn(3, 3),)
model = M()

exported_prog = torch.export.export(model, inputs)
#torch.export.save(exported_prog, "pt2e.pt2")

#saved_ep=torch.export.load("pt2e.pt2")
saved_ep=exported_prog
print("-------------------------------")
saved_ep.graph_module.print_readable()
print("-------------------------------")
print(saved_ep.graph_module.graph)


module_partitions = get_source_partitions(saved_ep.graph_module.graph, [torch.nn.Linear, torch.nn.ReLU])
print("==============================")
for module_or_fn_type, partitions in module_partitions.items():
    print("type")
    print(module_or_fn_type)
    for p in partitions:
        print("partition")
        print(p)
        for node in p.params:
            print("node")
            print(node)

```

test result:
```
-------------------------------
class GraphModule(torch.nn.Module):
    def forward(self, p_linear1_weight: "f32[3, 3]", p_linear1_bias: "f32[3]", p_linear2_weight: "f32[5, 3]", p_linear2_bias: "f32[5]", x: "f32[3, 3]"):
         # File: /home/yguo18/tmp/tmp/yjguo.testkit/fp8/source_partitions.py:15 in forward, code: x = self.linear1(x)
        linear: "f32[3, 3]" = torch.ops.aten.linear.default(x, p_linear1_weight, p_linear1_bias);  x = None

         # File: /home/yguo18/tmp/tmp/yjguo.testkit/fp8/source_partitions.py:16 in forward, code: x = self.linear1(x)
        linear_1: "f32[3, 3]" = torch.ops.aten.linear.default(linear, p_linear1_weight, p_linear1_bias);  linear = p_linear1_weight = p_linear1_bias = None

         # File: /home/yguo18/tmp/tmp/yjguo.testkit/fp8/source_partitions.py:17 in forward, code: x = self.relu(x)
        relu: "f32[3, 3]" = torch.ops.aten.relu.default(linear_1);  linear_1 = None

         # File: /home/yguo18/tmp/tmp/yjguo.testkit/fp8/source_partitions.py:18 in forward, code: x = self.linear2(x)
        linear_2: "f32[3, 5]" = torch.ops.aten.linear.default(relu, p_linear2_weight, p_linear2_bias);  relu = p_linear2_weight = p_linear2_bias = None
        return (linear_2,)

-------------------------------
graph():
    %p_linear1_weight : [num_users=2] = placeholder[target=p_linear1_weight]
    %p_linear1_bias : [num_users=2] = placeholder[target=p_linear1_bias]
    %p_linear2_weight : [num_users=1] = placeholder[target=p_linear2_weight]
    %p_linear2_bias : [num_users=1] = placeholder[target=p_linear2_bias]
    %x : [num_users=1] = placeholder[target=x]
    %linear : [num_users=1] = call_function[target=torch.ops.aten.linear.default](args = (%x, %p_linear1_weight, %p_linear1_bias), kwargs = {})
    %linear_1 : [num_users=1] = call_function[target=torch.ops.aten.linear.default](args = (%linear, %p_linear1_weight, %p_linear1_bias), kwargs = {})
    %relu : [num_users=1] = call_function[target=torch.ops.aten.relu.default](args = (%linear_1,), kwargs = {})
    %linear_2 : [num_users=1] = call_function[target=torch.ops.aten.linear.default](args = (%relu, %p_linear2_weight, %p_linear2_bias), kwargs = {})
    return (linear_2,)
==============================
type
<class 'torch.nn.modules.linear.Linear'>
partition
SourcePartition(nodes=[p_linear1_weight, p_linear1_bias, linear_1], source=<class 'torch.nn.modules.linear.Linear'>, input_nodes=[linear], output_nodes=[p_linear1_weight, linear_1, p_linear1_bias], params=[])
partition
SourcePartition(nodes=[p_linear2_weight, p_linear2_bias, linear_2], source=<class 'torch.nn.modules.linear.Linear'>, input_nodes=[relu], output_nodes=[linear_2], params=[])
partition
SourcePartition(nodes=[linear], source=<class 'torch.nn.modules.linear.Linear'>, input_nodes=[x, p_linear1_bias, p_linear1_weight], output_nodes=[linear], params=[])
type
<class 'torch.nn.modules.activation.ReLU'>
partition
SourcePartition(nodes=[relu], source=<class 'torch.nn.modules.activation.ReLU'>, input_nodes=[linear_1], output_nodes=[relu], params=[])

```
**Q1)** a tiny concern (https://github.com/pytorch/pytorch/pull/98628/files#r1825656920) about the output_nodes.
**Q2)** For "SourcePartition(nodes=[linear], source=<class 'torch.nn.modules.linear.Linear'>, input_nodes=[x, p_linear1_bias, p_linear1_weight], output_nodes=[linear], params=[])", why p_linear1_bias and p_linear1_weight are not in nodes, but in input_nodes.  It does not align with linear1 and linear2.

### **case 2:**   
save the PT2E model to disk, load it and then call get_source_partitions, 

test code:
```
import torch
from torch.fx.passes.utils.source_matcher_utils import (
                    SourcePartition,
                    get_source_partitions,
)

class M(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = torch.nn.Linear(3, 3)
        self.relu = torch.nn.ReLU()
        self.linear2 = torch.nn.Linear(3, 5)

    def forward(self, x):
        x = self.linear1(x)
        x = self.linear1(x)
        x = self.relu(x)
        x = self.linear2(x)
        return x

inputs = (torch.randn(3, 3),)
model = M()

exported_prog = torch.export.export(model, inputs)
torch.export.save(exported_prog, "pt2e.pt2")

saved_ep=torch.export.load("pt2e.pt2")
#saved_ep=exported_prog
print("-------------------------------")
saved_ep.graph_module.print_readable()
print("-------------------------------")
print(saved_ep.graph_module.graph)


module_partitions = get_source_partitions(saved_ep.graph_module.graph, [torch.nn.Linear, torch.nn.ReLU])
print("==============================")
for module_or_fn_type, partitions in module_partitions.items():
    print("type")
    print(module_or_fn_type)
    for p in partitions:
        print("partition")
        print(p)
        for node in p.params:
            print("node")
            print(node)

```

run result:
```
-------------------------------
class GraphModule(torch.nn.Module):
    def forward(self, p_linear1_weight: "f32[3, 3]", p_linear1_bias: "f32[3]", p_linear2_weight: "f32[5, 3]", p_linear2_bias: "f32[5]", x: "f32[3, 3]"):
         # File: /home/yguo18/tmp/tmp/yjguo.testkit/fp8/source_partitions.py:15 in forward, code: x = self.linear1(x)
        linear: "f32[3, 3]" = torch.ops.aten.linear.default(x, p_linear1_weight, bias = p_linear1_bias);  x = None

         # File: /home/yguo18/tmp/tmp/yjguo.testkit/fp8/source_partitions.py:16 in forward, code: x = self.linear1(x)
        linear_1: "f32[3, 3]" = torch.ops.aten.linear.default(linear, p_linear1_weight, bias = p_linear1_bias);  linear = p_linear1_weight = p_linear1_bias = None

         # File: /home/yguo18/tmp/tmp/yjguo.testkit/fp8/source_partitions.py:17 in forward, code: x = self.relu(x)
        relu: "f32[3, 3]" = torch.ops.aten.relu.default(linear_1);  linear_1 = None

         # File: /home/yguo18/tmp/tmp/yjguo.testkit/fp8/source_partitions.py:18 in forward, code: x = self.linear2(x)
        linear_2: "f32[3, 5]" = torch.ops.aten.linear.default(relu, p_linear2_weight, bias = p_linear2_bias);  relu = p_linear2_weight = p_linear2_bias = None
        return (linear_2,)

-------------------------------
graph():
    %p_linear1_weight : [num_users=2] = placeholder[target=p_linear1_weight]
    %p_linear1_bias : [num_users=2] = placeholder[target=p_linear1_bias]
    %p_linear2_weight : [num_users=1] = placeholder[target=p_linear2_weight]
    %p_linear2_bias : [num_users=1] = placeholder[target=p_linear2_bias]
    %x : [num_users=1] = placeholder[target=x]
    %linear : [num_users=1] = call_function[target=torch.ops.aten.linear.default](args = (%x, %p_linear1_weight), kwargs = {bias: %p_linear1_bias})
    %linear_1 : [num_users=1] = call_function[target=torch.ops.aten.linear.default](args = (%linear, %p_linear1_weight), kwargs = {bias: %p_linear1_bias})
    %relu : [num_users=1] = call_function[target=torch.ops.aten.relu.default](args = (%linear_1,), kwargs = {})
    %linear_2 : [num_users=1] = call_function[target=torch.ops.aten.linear.default](args = (%relu, %p_linear2_weight), kwargs = {bias: %p_linear2_bias})
    return (linear_2,)
==============================
type
<class 'torch.nn.modules.linear.Linear'>
partition
SourcePartition(nodes=[linear], source=<class 'torch.nn.modules.linear.Linear'>, input_nodes=[p_linear1_weight, x], output_nodes=[linear], params=[])
partition
SourcePartition(nodes=[linear_1], source=<class 'torch.nn.modules.linear.Linear'>, input_nodes=[p_linear1_weight, linear], output_nodes=[linear_1], params=[])
partition
SourcePartition(nodes=[linear_2], source=<class 'torch.nn.modules.linear.Linear'>, input_nodes=[relu, p_linear2_weight], output_nodes=[linear_2], params=[])
type
<class 'torch.nn.modules.activation.ReLU'>
partition
SourcePartition(nodes=[relu], source=<class 'torch.nn.modules.activation.ReLU'>, input_nodes=[linear_1], output_nodes=[relu], params=[])

```
**Q3)** why the source partition result is different than case 1?
**Q4)** why p_linear*_bias is not there

### Versions

pip list | grep torch
torch                    2.5.1

thanks


cc @jerryzh168 @jianyuh @raghuramank100 @jamesr66a @vkuzo @jgong5 @Xia-Weiwen @leslie-fang-intel @msaroufim @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4 @penguinwu