modify unflatten for vllm

[liangel-02]

Summary

We keep track of what has been processed by deleting processed keys from the original state dict. This is for the case when not all tensor attributes are available to us when loading, so we reconstruct the tensor subclasses with missing attributes in a later call.

We expect the state dict to be empty after loading is complete

Test
python test/prototype/safetensors/test_safetensors_support.py

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3297

• 📄 Preview Python docs built from this PR

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

⏳ No Failures, 1 Pending

As of commit f179ca5 with merge base e2aab90 ():
💚 Looks good so far! There are no failures yet. 💚

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

[jerryzh168]

I think current impl is a bit hacky, what we can do is:

unflatten_tensor_state_dict with incomplete state_dict:
   return a tuple of (unflattened_state_dict, unused part of incomplete_state_dict)
   
next call: combine the new state_dict input, with the previous unconsumed part

[jerryzh168]

looks good, thanks!