Ensure that tied parameter is children of module #3327
Conversation
|
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. |
| @@ -882,6 +883,48 @@ def test_get_balanced_memory(self): | |||
| max_memory = get_balanced_memory(model, max_memory={0: 0, "cpu": 100}) | |||
| assert {0: 0, "cpu": 100} == max_memory | |||
|
|
|||
| def test_get_module_size_with_ties(self): | |||
| # Create a model with a ModuleList containing more than 10 elements | |||
There was a problem hiding this comment.
Let's add a reference to issue #3308 here.
tests/test_modeling_utils.py
Outdated
| # Retrieve the tied parameters. Each group of tied parameters is returned | ||
| # by find_tied_parameters in alphabetical order | ||
| tied_parameters = find_tied_parameters(model) | ||
| assert tied_parameters == [sorted([f"{i}.weight" for i in range(num_layers)])] |
There was a problem hiding this comment.
For my understanding, are the checks starting from here really necessary to cover the bug fix? Or could the test be finished here?
There was a problem hiding this comment.
The call to find_tied_parameters could be omitted, as it is not really the function under test, and just set tied_parameters=[sorted([f"{i}.weight" for i in range(num_layers)])], thus removing the first assert. I included that call to emphasise the flow in infer_auto_device_map. I agree that it would be better to remove it, so the test only covers the bug fix (which is in get_module_size_with_ties). I guess that the rest of the test would be fine as it is, what do you think?
pablomlago
force-pushed
the
fix-get-modules-with-ties
branch
from
348c158 to
5ff6fb7
Compare
Ensure that tied parameters are assigned to their parent module in get_module_size_with_ties Fixes: huggingface#3308
pablomlago
force-pushed
the
fix-get-modules-with-ties
branch
from
5ff6fb7 to
21df343
Compare
Fixes #3308
This PR fixes the check in get_module_size_with_ties used to determine the module to which a tied parameter belongs. This is done by checking that the name of the tied parameter starts with module name plus ".".
The previous check (
n in tied_param) was not returning the expected results in models with a ModuleList, where the contained modules have a tied parameter. I added a unit test covering this situation.cc @BenjaminBossan