riaz.somc/transformers
1
0
Fork 0
mirror of https://github.com/huggingface/transformers.git synced 2025-08-01 02:31:11 +06:00
Code Issues Actions 3 Packages Projects Releases Wiki Activity

PegasusX add _no_split_modules (#25933)

Browse Source 
* no_split_modules

* no_split_modules

* inputs_embeds+pos same device

* update _no_split_modules

* update _no_split_modules
This commit is contained in:
andreeahedes 2023-09-05 16:34:34 +01:00 committed by GitHub
parent 70a98024b1
commit da1af21dbb
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
src/transformers/models/pegasus_x/modeling_pegasus_x.py
View File

@ -769,6 +769,7 @@ class PegasusXPreTrainedModel(PreTrainedModel):
config_class = PegasusXConfig
base_model_prefix = "model"
supports_gradient_checkpointing = True
_no_split_modules = [r"PegasusXEncoderLayer", r"PegasusXDecoderLayer"]
def _init_weights(self, module):
std = self.config.init_std
@ -1299,6 +1300,8 @@ class PegasusXDecoder(PegasusXPreTrainedModel):
# embed positions
positions = self.embed_positions(inputs_embeds, past_key_values_length)
positions = positions.to(inputs_embeds.device)
hidden_states = inputs_embeds + positions
hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)