[torchao] safetensors

[liangel-02]

Summary
updating safetensors logic to support all tensor subclasses generically

Test
verified integration with vllm
unit tests pass: python tests/quantization/torchao_integration/test_torchao.py -k TorchAoSafeSerializationTest

[liangel-02]

cc @SunMarc @MekkCyber

[MekkCyber]

Thanks a lot for working on this

[MekkCyber]

why do we need this ? could you explain a bit more why we need to change the regex handling here and in the other parts of the code?

[liangel-02]

for the torchao WeightConverter, the original code has hardcoded tensor data components (ie.
weight_qdata, weight_scale) mapped to the consolidated weight. however, these components change depending on the config used, with some also being optional. for max generality, i wanted to use wildcard matching (*_weight_* --> *weight*).

but with the current regex handling, the code was renaming the key with the literal "*weight*" rather than matching the regex. the changes i have extracts the prefix and uses that for source/target. lmk if this is an ok approach/if theres a better solution

[liangel-02]

@MekkCyber to avoid adding torchao specific logic to the core model loader and unblock landing, i reverted the changes to the regex logic and added all other tensor data names. in the future, maybe we can implement a strategy that generalizes this matching in a better way?

cc @jerryzh168

[jerryzh168]

for this one, maybe just change it to (f"_{x} for x in self.full_ao_keys)" to be safe, at least with this we are confident it's correct. if there are no better general way to detect safetensors

[liangel-02]

i think this is achieving the same thing as L237, i think we only need one check

[MekkCyber]

Very nice alternative solution ! Thanks a lot

[MekkCyber]

let's do an if/else and raise an error to follow the same pattern we use in transformers, but i'm not sure if it's necessary since if we have the suffix in input_dict it means we collected at least a tensor for that suffix

[liangel-02]

will add an if/else! we wanted to add this assert since there also shouldn't be more than one tensor for a component (qdata, scale, etc)

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: torchao_integration

[liangel-02]

@MekkCyber @SunMarc can you help me merge this, thanks!