[torchao] fix safetensors for sharding

[liangel-02]

Summary

Previously, our safetensors implementation for tensor subclasses didn't consider the possibility of different tensor attributes for one instance of a tensor subclass (ie qdata or scale for Float8Tensor) being sharded to different files.

This PR handles that case by:

1. Defining the state dict outside of the loop that iterates through shard files
2. After processing a full tensor subclass in unflatten_tensor_state_dict, we delete the corresponding tensor attributes (ie qdata, scale) from the original state dict passed in
3. Whatever is remaining is saved in the state dict for the next iteration

Testing
pytest tests/quantization/test_torchao.py::test_safetensors_model_loading_with_params -s

[heheda12345]

do we need to skip now that torch2.9 is released?

[liangel-02]

i think we still need to skip since this test depends on changes in torchao main cc @jerryzh168

[jerryzh168]

yeah we need a separate update to upgrade the torchao version used in testing in vllm:

vllm/.buildkite/test-amd.yaml

Line 616 in da786e3

- uv pip install --system torchao==0.13.0

and make sure the tests are passing

[heheda12345]

LGTM!

[liangel-02]

@heheda12345 can you pls merge? thanks!

[heheda12345]

Of course! I'm retrying failed test.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @liangel-02.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[liangel-02]

@heheda12345 a rebase seems to have fixed ci, can u help me merge ty!

[heheda12345]

Merged. Thanks!