Change `wte` to use shared embedding

modeling_mpt.py

@@ -12,6 +12,7 @@ from transformers import PreTrainedModel, PreTrainedTokenizer, PreTrainedTokeniz
 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
 from .attention import attn_bias_shape, build_attn_bias
 from .blocks import MPTBlock
+from .custom_embedding import SharedEmbedding
 from .norm import NORM_CLASS_REGISTRY
 from .configuration_mpt import MPTConfig
 from .adapt_tokenizer import AutoTokenizerForMOD, adapt_tokenizer_for_denoising
@@ -44,7 +45,7 @@ class MPTModel(MPTPreTrainedModel):
             raise NotImplementedError(f'Requested norm type ({config.norm_type}) is not implemented within this repo (Options: {norm_options}).')
         norm_class = NORM_CLASS_REGISTRY[config.norm_type.lower()]
         self.embedding_fraction = config.embedding_fraction
-        self.wte = nn.Embedding(config.vocab_size, config.d_model, device=config.init_device)
+        self.wte = SharedEmbedding(config.vocab_size, config.d_model, device=config.init_device)
         if not self.alibi:
             self.wpe = nn.Embedding(config.max_seq_len, config.d_model, device=config.init_device)
         self.emb_drop = nn.Dropout(config.emb_pdrop)
@@ -252,7 +253,7 @@ class MPTForCausalLM(MPTPreTrainedModel):
         return_dict = return_dict if return_dict is not None else self.config.return_dict
         use_cache = use_cache if use_cache is not None else self.config.use_cache
         outputs = self.transformer(input_ids=input_ids, past_key_values=past_key_values, attention_mask=attention_mask, prefix_mask=prefix_mask, sequence_id=sequence_id, return_dict=return_dict, output_attentions=output_attentions, output_hidden_states=output_hidden_states, use_cache=use_cache)
-        logits = F.linear(outputs.last_hidden_state.to(self.transformer.wte.weight.device), self.transformer.wte.weight)
+        logits = self.transformer.wte(outputs.last_hidden_state.to(self.transformer.wte.weight.device), True)
         if self.logit_scale is not None:
             if self.logit_scale == 0:
                 warnings.warn(f'Multiplying logits by self.logit_scale={self.logit_scale!r}. This will produce uniform (uninformative) outputs.')