Repository: pytorch
Issue: Fx graph recompilation along with partitioner produces different node order even though graph operations are same
Labels: high priority,triage review,triaged,oncall: fx

### 🐛 Describe the bug

![gm_node_order_1](https://github.com/pytorch/pytorch/assets/39515476/c56e7928-68d3-4df0-9367-09487b12ea51)
```
import torch
import torch.nn as nn
aten = torch.ops.aten
from torch._functorch.aot_autograd import aot_module_simplified
import numpy as np
import random

from typing import Mapping
from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner
from torch.fx.passes.operator_support import OperatorSupport

class DevOperatorSupport(OperatorSupport):
    def is_node_supported(self, submodules: Mapping[str, torch.nn.Module], node: torch.fx.Node) -> bool:
        return True


class DevPartitioner(CapabilityBasedPartitioner):
    def __init__(self, graph_module: torch.fx.GraphModule):
        super().__init__(
            graph_module,
            DevOperatorSupport(),
            allows_single_node_partition=True,
        )
def my_aot_compiler(gm, example_inputs):
    def my_compiler(gm, example_inputs):
        # print(gm.code)
        # print(gm.graph)
        import copy
        from torch._functorch.compile_utils import strip_overloads
        from torch._functorch.compilers import _disable_jit_autocast

        module = copy.deepcopy(gm)

        with _disable_jit_autocast():
            strip_overloads(module)

            for node in module.graph.nodes:
                new_kwargs = {}
                for k, v in node.kwargs.items():
                    if isinstance(v, torch.device):
                        v = v.type
                    new_kwargs[k] = v
                node.kwargs = new_kwargs

            DevPartitioner(module).propose_partitions()
            module.graph.lint()
            module.recompile()

            from collections import OrderedDict
            module._forward_hooks = OrderedDict()
            module._forward_pre_hooks = OrderedDict()
            print("Module is \n", module.print_readable())
            f = torch.jit.script(module)
            print("Converted Graph is \n", f.graph)
            #DevPartitioner(gm).propose_partitions()
        return gm.forward


    # Invoke AOTAutograd
    return aot_module_simplified(
        gm,
        example_inputs,
        fw_compiler=my_compiler
    )

class Net(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Sequential(torch.nn.Conv2d(1, 6, 5, 1, 0),)

    def forward(self, input):
        result = out = self.layer(input)
        return result

def test_dynamic_shape_topk_cpu():
    print("Starting the test.................")
    sizes = [
        (8, 1, 16, 16),
    ]

    model_cpu = Net().to('cpu')

    compiled_function_training = torch.compile(model_cpu, backend=my_aot_compiler, dynamic=True)

    with torch.no_grad():
        with torch.autocast(device_type="cpu", dtype=torch.bfloat16, enabled=True):
            t = torch.rand(8, 1, 16, 16)
            t = t.to(torch.bfloat16)
            result_compile_train = compiled_function_training(t)

test_dynamic_shape_topk_cpu()
```

### Versions

```
Collecting environment information...
PyTorch version: 2.1.0a0+gite999726
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.3 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: 14.0.0-1ubuntu1.1
CMake version: version 3.27.9
Libc version: glibc-2.35

Python version: 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-89-generic-x86_64-with-glibc2.35
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Address sizes:                      43 bits physical, 48 bits virtual
Byte Order:                         Little Endian
CPU(s):                             12
On-line CPU(s) list:                0-11
Vendor ID:                          GenuineIntel
Model name:                         Intel(R) Xeon(R) Gold 5220R CPU @ 2.20GHz
CPU family:                         6
Model:                              85
Thread(s) per core:                 1
Core(s) per socket:                 6
Socket(s):                          2
Stepping:                           0
BogoMIPS:                           4389.68
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon nopl xtopology tsc_reliable nonstop_tsc cpuid tsc_known_freq pni pclmulqdq vmx ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single pti ssbd ibrs ibpb stibp tpr_shadow vnmi ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 invpcid avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xsaves arat pku ospke md_clear flush_l1d arch_capabilities
Virtualization:                     VT-x
Hypervisor vendor:                  VMware
Virtualization type:                full
L1d cache:                          384 KiB (12 instances)
L1i cache:                          384 KiB (12 instances)
L2 cache:                           12 MiB (12 instances)
L3 cache:                           71.5 MiB (2 instances)
NUMA node(s):                       1
NUMA node0 CPU(s):                  0-11
Vulnerability Gather data sampling: Unknown: Dependent on hypervisor status
Vulnerability Itlb multihit:        KVM: Mitigation: VMX disabled
Vulnerability L1tf:                 Mitigation; PTE Inversion; VMX flush not necessary, SMT disabled
Vulnerability Mds:                  Mitigation; Clear CPU buffers; SMT Host state unknown
Vulnerability Meltdown:             Mitigation; PTI
Vulnerability Mmio stale data:      Mitigation; Clear CPU buffers; SMT Host state unknown
Vulnerability Retbleed:             Mitigation; IBRS
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; IBRS, IBPB conditional, STIBP disabled, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] numpy==1.23.5
[pip3] torch==2.1.0a0+gite999726
```

cc @ezyang @gchanan @zou3519 @kadeng @SherlockNoMad @EikanWang @jgong5 @wenzhe-nrv