Switch to new ao quant api for 8da4w

[jackzhxng]

Summary

Closes #8422

Test plan

[PLEASE REMOVE] How did you test this PR? Please write down any manual commands you used and note down tests that you have written if applicable.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/8501

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 285d20b with merge base 6cb5c1a ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[jerryzh168]

thanks, please make sure the lowering still works as well

[mcr229]

so i believe the qo quant api introduces quant_affine nodes, which are decomposed. We got around this by using to_edge_transform_and_lower path to preserve them from decomposition, but I think that would also require us to modify the llama_export path to use to_edge_transform_and_lower as well.

[jackzhxng]

@mcr229 taking a stab at that here #8624

[jackzhxng]

^Above has been closed and xnnpack is now using to_edge_transform_and_lower, giving CI another crack at this

[mcr229]

looks like this is running into issues with loading the quantized checkpoints because aten._copy is unimplemented? seems sus, let me know if you need help from me for any part of enabling this.

[jackzhxng]

Yeah this is weird, our one test for quantized embeddings is erroring, presumably because the embedding quantization happens after the linear quantization which is now using the API

[jerryzh168]

please check out https://pytorch.org/ao/stable/serialization.html to see how serialization works in torchao, copy_ a torchao quantized tensor to a normal tensor is indeed not supported, we typically use load_state_dict(...., assign=True)

also embedding quantization is not yet fully supported in torchao I think, but we may add this in H1

[jackzhxng]

@jerryzh168 this error is being thrown during the load_state_dict here we are already doing which loads the quantized embedding weights - https://github.com/pytorch/executorch/blob/main/examples/models/llama/source_transformation/quantize.py#L666, any idea where this copy_ op is coming from?

[jackzhxng]

Update - should be resolved. This error was happening since the quantize embedding code was loading the model's state dict which contains ao-quantized tensors from the linear transform before which now uses quantize_. Need the assign=True to do the deserialization process properly. Thanks @andrewor14 @jerryzh168 @mcr229

[jackzhxng]

Actually I should have imported and run internal tests before merging, if diff train gets blocked I'll forward fix

[jerryzh168]

is torch_dtype not applied here? should it be applied to model?

[jackzhxng]

Yeah I think it would good to ensure the dtype here by applying it to the model in general, but the model is already in fp32 for when the test passes before this PR and when it fails after this PR