[Inductor][Compile]torch.compile remembers the old weight for evaluation if grad is disabled and it runs on CPU
Labels: triaged, oncall: pt2, module: inductor

### 🐛 Describe the bug

```python
import torch
from torch import nn
import copy

class BackboneModel(nn.Module):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.conv1 = nn.Conv2d(16, 16, 6)
        self.bn1 = nn.BatchNorm2d(16)
    def forward(self, x):
        x = self.bn1(self.conv1(x))
        return x

model = BackboneModel().eval()
optim = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

old_model = copy.deepcopy(model) # remember the old model
opt_model = torch.compile(model) # compile the model

a = torch.rand(64, 16, 32, 32)

with torch.no_grad():
    model.eval()
    output1 = model(a)
    output2 = opt_model(a)

print("diff of eager model and optimized model at initialization:")
print((output1 - output2).abs().max().item()) # prints 4.172325134277344e-07 on my computer

# train the model and update the weight
model.train()
for i in range(10):
    optim.zero_grad()
    output = opt_model(a)
    output.sum().backward()
    optim.step()

with torch.no_grad():
    model.eval()
    output1 = model(a)
    output2 = opt_model(a)
    output3 = old_model(a)

print("diff of eager model and optimized model after training:")
print((output1 - output2).abs().max().item()) # prints 19307.39453125 on my computer

print("diff of old model and optimized model after training:")
# the opt_model remembers the old model parameters, and don't use the new weight!
print((output3 - output2).abs().max().item()) # prints 4.172325134277344e-07 on my computer
```
The reason:

When grad is disabled and conv-bn modules are in eval mode, the `fuse_conv_bn` optimization is turned on:

https://github.com/pytorch/pytorch/blob/09f7cb0eaf02d69ab4c788b03345198f5c869718/torch/_inductor/fx_passes/pre_grad.py#L97

The optimization fuses parameters directly, creating a computation graph with **a new conv** whose parameters are computed from the conv and bn parameters **at the time of compiling the model**.

Consider the procedure of eval-train-eval:
1. When we evaluate the model at the first time, the newly created conv in the compiled graph remembers the old parameter values;
2. Then when we train the model, the parameters' value changed, but the newly created conv is not updated.
3. When we evaluate the model again, the guarding conditions for the eval graph remains satisfied, so this doesn't trigger re-compilation. The eval graph constructed in Step 1 is used directly, with the obsolete values.

The bug affects all the training loop who uses `torch.compile` and evaluates the model with `torch.no_grad()`.

This bug is very difficult to find and check, as only consecutive conv-bn values in a whole neural network are affected. The eval graph runs with some updated parameters and some obselete parameters. Only when I dive into the details and shrink the network to have only one pair of conv-bn, does the bug appear clearly.

A possible fix is to fuse computation rather than parameters for conv-bn, as described in https://github.com/pytorch/pytorch/pull/108757 : 

- When we find a pair of conv-bn usage, as long as the bn is in eval mode, we change the computation of `normalize (weight conv input)` to `(normalize (weight)) conv input`. These two are mathematically equivalent, and does not touch the parameters.
- If users explicit specify something like `TORCHINDUCTOR_FREEZING=1`, we know the parameters cannot be updated. Then the constant folding pass can automatically fold the `normalize (weight)` into a baked `weight`.

This solution is sound, and can also deal with training conv-bn when bn is kept in eval mode (which reduces a lot of training memory footprint).

### Versions

I'm testing on `2.1.0.dev20230904`, as well as nightly `2.2` built from main.

CPU versions are affected. Fortunately, GPU versions are not affected. The `fuse_conv_bn` happens if no grad and cpu and conv-bn eval are all true:

https://github.com/pytorch/pytorch/blob/09f7cb0eaf02d69ab4c788b03345198f5c869718/torch/_inductor/fx_passes/pre_grad.py#L77

```python
def fuse_fx(gm: torch.fx.GraphModule, example_inputs):
    is_cpu = is_cpu_device(example_inputs)

    fake_mode = detect_fake_mode(example_inputs)

    gm = sink_cat_after_pointwise(gm)
    if config.permute_fusion and not is_cpu:
        # For linear permute fusion, we need to check input info to identify
        # and perform proper permutation/transpose
        ShapeProp(gm, fake_mode=fake_mode).propagate(*example_inputs)
        gm = linear_permute_fusion(gm)
        gm = permute_linear_fusion(gm)
        gm = permute_matmul_fusion(gm)

    # make sure the autograd is disabled.
    if torch.is_grad_enabled():
        return gm
    if not is_cpu:
        return gm
    gm = remove_identity(gm)
    gm = fuse_conv_bn(gm)
    return gm

```

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305 @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @ngimel @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov