Skip to content

don't initialize the output embeddings if we're going to tie them to input embeddings #28192

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ArthurZucker merged 2 commits into from
Jan 31, 2024
+39 −0
Merged

don't initialize the output embeddings if we're going to tie them to input embeddings #28192

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0c0efed
test that tied output embeddings aren't initialized on load
tom-p-reichel Jan 23, 2024
28fe452
don't initialize the output embeddings if we're going to tie them to …
tom-p-reichel Dec 22, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions src/transformers/modeling_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3746,6 +3746,11 @@ def _fix_key(key):
else:
_loaded_keys = loaded_keys
not_initialized_submodules = set_initialized_submodules(model, _loaded_keys)
# if we're about to tie the output embeds to the input embeds we don't need to init them
if hasattr(model.config, "tie_word_embeddings") and model.config.tie_word_embeddings:
output_embeddings = model.get_output_embeddings()
if output_embeddings is not None:
output_embeddings._is_hf_initialized = True
Comment on lines 3749 to 3753
Copy link
Copy Markdown
Collaborator

@ArthurZucker ArthurZucker Jan 5, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Very nice!
I think we can also add something like

transformers/src/transformers/modeling_utils.py

Line 1580 in 90224dd

def tie_weights(self):

for

        if getattr(self.config, "is_encoder_decoder", False) and getattr(self.config, "tie_encoder_decoder", False):
            if hasattr(self, self.base_model_prefix):
                self = getattr(self, self.base_model_prefix)
            self._tie_encoder_decoder_weights(self.encoder, self.decoder, self.base_model_prefix)

which can be tested on MarianMT models but should come in another PR!

else:
not_initialized_submodules = dict(model.named_modules())
# This will only initialize submodules that are not marked as initialized by the line above.
Expand Down
34 changes: 34 additions & 0 deletions tests/test_modeling_common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -483,6 +483,40 @@ def _init_weights(self, module):
max_diff = torch.max(torch.abs(model_slow_init.state_dict()[key] - model_fast_init.state_dict()[key]))
self.assertLessEqual(max_diff.item(), 1e-3, msg=f"{key} not identical")

def test_fast_init_tied_embeddings(self):
class MyClass(PreTrainedModel):
config_class = PretrainedConfig
_tied_weights_keys = ["output_embeddings.weight"]

def __init__(self, config=None):
super().__init__(config if config is not None else PretrainedConfig())
self.input_embeddings = nn.Embedding(10, 10)
self.output_embeddings = nn.Linear(10, 10, bias=False)
self.tie_weights()

def get_output_embeddings(self):
return self.output_embeddings

def set_output_embeddings(self, output_embeddings):
self.output_embeddings = output_embeddings

def get_input_embeddings(self):
return self.input_embeddings

def set_input_embeddings(self, input_embeddings):
self.input_embeddings = input_embeddings

def _init_weights(self, module):
if module is self.output_embeddings:
raise ValueError("unnecessarily initialized tied output embedding!")

model = MyClass()

with tempfile.TemporaryDirectory() as tmpdirname:
model.save_pretrained(tmpdirname)
# throws if it initializes the tied output_embeddings
MyClass.from_pretrained(tmpdirname)

def test_save_load_fast_init_to_base(self):
config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
if config.__class__ not in MODEL_MAPPING:
Expand Down