Johnblick187 commited on
Commit
c06d951
·
verified ·
1 Parent(s): 9d345fa

Update modeling_smartcoder_moe.py

Browse files
Files changed (1) hide show
  1. modeling_smartcoder_moe.py +8 -0
modeling_smartcoder_moe.py CHANGED
@@ -245,6 +245,14 @@ class SmartCoderMoEForCausalLM(PreTrainedModel, GenerationMixin):
245
  self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
246
  self.post_init()
247
 
 
 
 
 
 
 
 
 
248
  def get_input_embeddings(self): return self.model.embed_tokens
249
  def get_output_embeddings(self): return self.lm_head
250
 
 
245
  self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
246
  self.post_init()
247
 
248
+ def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
249
+ remapped = {}
250
+ for k, v in state_dict.items():
251
+ k = k.replace('experts_fc.weight', 'experts_fc')
252
+ k = k.replace('experts_proj.weight', 'experts_proj')
253
+ remapped[k] = v
254
+ super()._load_from_state_dict(remapped, prefix, *args, **kwargs)
255
+
256
  def get_input_embeddings(self): return self.model.embed_tokens
257
  def get_output_embeddings(self): return self.lm_head
258