Correct loading of models with shared tensors when using accelerator.load_state() #2875
Conversation
…lerator.load_state()
|
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. |
| models[i].load_state_dict(state_dict, **load_model_func_kwargs) | ||
| model.load_state_dict(state_dict, **load_model_func_kwargs) | ||
| logger.info("All model weights loaded successfully") |
There was a problem hiding this comment.
Any particular reason for this change? I'd expect only the prior to be modified.
There was a problem hiding this comment.
In the if statement, he's loading the safetensors model directly whereas before, we were only getting the state dict.
There was a problem hiding this comment.
load_model does both: it loads the file and uses it to populate the state_dict. Previously, each branch of the if-condition only loaded the file and after the if-condition, the model would load the state dict. Since load_model does both, I indented the statement on line 204 to become part of the else-clause. This becomes clearer when you have a look at the complete surroundings of the changes instead of only the affected lines.
There was a problem hiding this comment.
The other change in this line (aside from the indent) namely using model instead of models[i] is mostly cosmetic. My linter was complaining that the enumerate call defines model but it's never used.
There was a problem hiding this comment.
Thanks for the change and spotting the issue ! Could you add a test with a model with tied weight ? You can use the following test for reference : test_save_load_model
|
Yes, I'll have a look into it. |
|
You can verify that the shared weights are implemented correctly by checking the output. safetensors warns you about that fact. |
|
Just did. This is the expected error message I get when reverting my changes. |
tests/test_accelerator.py
Outdated
| # need to add this for compliance with other methods | ||
| self.weight = self.linear1.weight | ||
| self.bias = self.linear1.bias |
There was a problem hiding this comment.
Do we really need that ? where does it fail ?
There was a problem hiding this comment.
It used to fail previously. You're right. This part can be safely removed.
I was only expecting linear2.weight and linear2.bias to be missing. Maybe this is due to self.weight = self.linear1.weight
self.bias = self.linear1.bias |
After removing the unnecessary bits, it correctly only throws an error for the weights and bias of the 2nd linear layer. |
|
Nice ! Could you just fix the quality issue ( |
What does this PR do?
I would run into problems with PyTorch's load_state_dict complaining about missing keys. These keys belonged to shared tensors. These shared keys are intentionally omitted by the safetensors library. To load a model correctly, one has to use safetensor's load_model function instead of the default load_state_dict function (described here). This was previously not done when using the load_state function of the Accelerator.
Fixes # (issue)
I think this issue might be relevant as they also report problems when loading with accelerator.load_state.
#2155
Before submitting
Pull Request section?
to it if that's the case.
documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.