Torchao weights only compability

[SunMarc]

What does this PR do ?

This PR makes torchao serialized model loadable with weights_only=True which is the default. Otherwise, you need to set weights_only=False which is not recommended.

cc @jerryzh168 cc @MekkCyber

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[jerryzh168]

if we do import torchao, I think we should get everything here (classes etc. being added to safeglobals)? otherwise we'd need to fix torchao

[SunMarc]

I'm using torchao 0.5.0 and it's not working on my side. I can try with the latest tomorrow !

[jerryzh168]

I see, it's not expected I think, I think it should be fixed in torchao side, I feel 0.5 should have this functionality already actually. if you can have a standalone repro that will be very helpful for us. I remember I have tested in https://huggingface.co/docs/transformers/main/en/quantization/torchao

[SunMarc]

Actually, we ran into this issue with @MekkCyber on the example you shared in the docs.
Here's a the reproducer, let us know if you also have this issue :

from transformers import TorchAoConfig, AutoTokenizer, AutoModelForCausalLM
import torch

model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
quant_config = TorchAoConfig("int4_weight_only", group_size=32)
quantized_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="cuda:0",
    quantization_config=quant_config,
)
output_dir = "llama3-8b-int4wo-128"
quantized_model.save_pretrained(output_dir, safe_serialization=False)

loaded_quantized_model = AutoModelForCausalLM.from_pretrained(output_dir, device_map="cuda:0")

[jerryzh168]

OK will test and report back

[MekkCyber]

Thanks for this PR @SunMarc, really helpful !