fix(nn/models): initializing base transformer from a custom config (#149)

Author: maxreciprocate

trlx/trainer/nn/ilql_models.py

@@ -197,14 +197,14 @@ def __init__(
                 _hfconfig = transformers.deepspeed.HfDeepSpeedConfig(  # noqa: F841
                     config_path
                 )
+
         if isinstance(config, str):
             self.config = transformers.AutoConfig.from_pretrained(config)
+            self.base_model = transformers.AutoModelForCausalLM.from_pretrained(config)
         else:
             self.config = config
+            self.base_model = transformers.AutoModelForCausalLM.from_config(config)
 
-        self.base_model = transformers.AutoModelForCausalLM.from_pretrained(
-            self.config.name_or_path,
-        )
         self.base_model.transformer = hf_get_causal_base_model(self.base_model)
         self.base_model.lm_head = hf_get_lm_head(self.base_model)
         freeze_bottom_causal_layers(self.base_model, num_layers_unfrozen)

trlx/trainer/nn/ppo_models.py

@@ -232,11 +232,11 @@ def __init__(self, config: Union[transformers.PretrainedConfig, str]):
         super().__init__()
         if isinstance(config, str):
             self.config = transformers.AutoConfig.from_pretrained(config)
+            self.base_model = transformers.AutoModelForCausalLM.from_pretrained(config)
         else:
             self.config = config
-        self.base_model = transformers.AutoModelForCausalLM.from_pretrained(
-            self.config.name_or_path
-        )
+            self.base_model = transformers.AutoModelForCausalLM.from_config(config)
+
         self.base_model.transformer = hf_get_causal_base_model(self.base_model)
         self.base_model.lm_head = hf_get_lm_head(self.base_model)
         self.v_head = make_head(hf_get_hidden_size(self.config), 1)
@@ -304,13 +304,14 @@ def __init__(
         num_layers_unfrozen: int = -1,
     ):
         super().__init__()
+
         if isinstance(config, str):
             self.config = transformers.AutoConfig.from_pretrained(config)
+            self.base_model = transformers.AutoModelForCausalLM.from_pretrained(config)
         else:
             self.config = config
-        self.base_model = transformers.AutoModelForCausalLM.from_pretrained(
-            self.config.name_or_path
-        )
+            self.base_model = transformers.AutoModelForCausalLM.from_config(config)
+
         self.base_model.transformer = hf_get_causal_base_model(self.base_model)
         self.base_model.lm_head = hf_get_lm_head(self.base_model)
         self.v_head = make_head(hf_get_hidden_size(self.config), 1)