Support device_map="auto" when loading

#23
by shijie-wu - opened
Files changed (1) hide show
  1. modeling_mpt.py +1 -0
modeling_mpt.py CHANGED
@@ -23,6 +23,7 @@ Tokenizer = Union[PreTrainedTokenizer, PreTrainedTokenizerFast]
23
  class MPTPreTrainedModel(PreTrainedModel):
24
  config_class = MPTConfig
25
  base_model_prefix = 'model'
 
26
 
27
  class MPTModel(MPTPreTrainedModel):
28
 
 
23
  class MPTPreTrainedModel(PreTrainedModel):
24
  config_class = MPTConfig
25
  base_model_prefix = 'model'
26
+ _no_split_modules = ["MPTBlock"]
27
 
28
  class MPTModel(MPTPreTrainedModel):
29