fix cast output layer

wanli2014 · Jun 28, 2023 · a136eb4 · a136eb4
1 parent 7b67c38
commit a136eb4
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/src/utils/common.py b/src/utils/common.py
@@ -233,13 +233,16 @@ def get_input_embeddings(self):
         model.get_input_embeddings = MethodType(get_input_embeddings, model)
         model.lm_head = model.transformer.output_layer
         tokenizer.eos_token = "</s>"
-        output_embedding_layer_name = "transformer.output_layer"
+        output_embedding_base_layer = model.transformer
+        output_embedding_layer_name = "output_layer"
     else:
+        output_embedding_base_layer = model
         output_embedding_layer_name = "lm_head"
 
     model = prepare_model_for_training(
         model,
         finetuning_args.finetuning_type,
+        output_embedding_base_layer,
         output_embedding_layer_name
     ) if is_trainable else model
     model = init_adapter(model, model_args, finetuning_args, is_trainable)

diff --git a/src/utils/other.py b/src/utils/other.py
@@ -73,6 +73,7 @@ def get_logits_processor() -> LogitsProcessorList:
 def prepare_model_for_training(
         model: PreTrainedModel,
         finetuning_type: str,
+        output_embedding_base_layer: torch.nn.Module,
         output_embedding_layer_name: Optional[str] = "lm_head",
         use_gradient_checkpointing: Optional[bool] = True,
         layer_norm_names: Optional[List[str]] = ["layernorm"] # for chatglm setting
@@ -87,16 +88,16 @@ def prepare_model_for_training(
         model.gradient_checkpointing_enable()
         model.config.use_cache = False # turn off when gradient checkpointing is enabled
 
-    if finetuning_type != "full" and hasattr(model, output_embedding_layer_name):
-        output_embedding_layer: torch.nn.Linear = getattr(model, output_embedding_layer_name)
+    if finetuning_type != "full" and hasattr(output_embedding_base_layer, output_embedding_layer_name):
+        output_embedding_layer = getattr(output_embedding_base_layer, output_embedding_layer_name)
         input_dtype = output_embedding_layer.weight.dtype
 
         class CastOutputToFloat(torch.nn.Sequential):
 
             def forward(self, x: torch.Tensor) -> torch.Tensor:
                 return super().forward(x.to(input_dtype)).to(torch.float32)
 
-        setattr(model, output_embedding_layer_name, CastOutputToFloat(output_embedding_layer))
+        setattr(output_embedding_base_layer, output_embedding_layer_name, CastOutputToFloat(output_embedding_layer))
 
     return model