Repository: pytorch
Issue: When loading a model trained with QAT `(qat.pth)` using `model.load_state_dict`,  an error occurs: `RuntimeError: Error(s) in loading state_dict for GraphModule: Missing key(s), Unexpected key(s) in state_dict.`
Labels: awaiting response (this tag is deprecated),needs reproduction,oncall: export

### 🐛 Describe the bug

- I'm reporting this issue due to errors related to capture_pre_autograd_graph and torch.compile in QAT.
- Note: Apologies if there are any misunderstandings.

- Based on the following tutorial, I implemented QAT on my custom model (let's call it `MMM`) and saved the trained model `qat.pth` using the `state_dict` format. Observing the training loss, I confirmed that the model is learning properly.
- Tutorial: [Quantization Aware Training with PT2E](https://pytorch.org/tutorials/prototype/pt2e_quant_qat.html)

■Excerpt from train.py (・・・ indicates omissions)
...
model = MMM().to(device)
...
example_inputs = (torch.randn(・・・).to(device),)  
dynamic_shapes = tuple(
    {0: torch.export.Dim("dim", min=1, max=128)} if i == 0 else None
    for i in range(len(example_inputs))
)
exported_model = torch.export.export_for_training(
    model, example_inputs, dynamic_shapes=dynamic_shapes
).module()
quantizer = XNNPACKQuantizer()
quantizer.set_global(get_symmetric_quantization_config(is_qat=True))
print("Model has been annotated for quantization.")
model = prepare_qat_pt2e(exported_model, quantizer)
model.to(device)
print("Model is prepared for QAT.")
...
torch.ao.quantization.move_exported_model_to_train(model)
...
torch.save(model.state_dict(), os.path.join(results_dir, "qat.pth")) 

- Afterwards, when I tried to evaluate the performance using the trained model "qat.pth" by loading it with load_state_dict, I encountered an error stating that the keys do not match. Looking at the error, the key names have slightly changed (some . have become _, etc.), and the existence of parameters like activation_post_process (probably for PTQ) is causing an error where the model cannot be loaded.
**- How can I resolve this?**

■Excerpt from valid.py
…
model = MMM().to(device)
example_inputs = (torch.randn(・・・).to(device),)
exported_model = capture_pre_autograd_graph(model, example_inputs)
quantizer = XNNPACKQuantizer()
quantizer.set_global(get_symmetric_quantization_config(is_qat=True))
model = prepare_qat_pt2e(exported_model, quantizer)
model.load_state_dict(torch.load("pth_path"))
model = convert_pt2e(model)
torch.ao.quantization.move_exported_model_to_eval(model)


### Error logs

  File "・・・/valid.py", line 53, in main
    model.load_state_dict(torch.load("pth_path"))
  File "・・・/miniconda3/envs/pytorch2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 2584, in load_state_dict
    raise RuntimeError(
RuntimeError: Error(s) in loading state_dict for GraphModule:

    Missing key(s) in state_dict:
        "XXX_0_weight", "XXX_0_bias", ・・・,
        "YYY_0_weight", "YYY_0_bias",
        "XXX_1_num_batches_tracked", "XXX_1_running_mean", "XXX_1_running_var", ・・・,

    Unexpected key(s) in state_dict:
        "XXX.0.weight", "XXX.0.bias", ・・・,
        "YYY.0.weight", "YYY.0.bias",
        "XXX.1.num_batches_tracked", "XXX.1.running_mean", "XXX.1.running_var", ・・・,
        "XXX.3.weight", "XXX.3.bias", ・・・,
        "activation_post_process_Z.fake_quant_enabled", "activation_post_process_Z.observer_enabled", "activation_post_process_Z.scale", "activation_post_process_Z.zero_point", "activation_post_process_Z.activation_post_process.eps", "activation_post_process_Z.activation_post_process.min_val", "activation_post_process_Z.activation_post_process.max_val", ・・・


### Versions

Windows WSL
Miniconda
PyTorch 2.5.1
Python 3.10

cc @chauhang @penguinwu @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4