support falcon model #72

2023-07-05 15:00:06 +08:00
parent 966b5c70fc
commit c136f362c1
5 changed files with 46 additions and 29 deletions
--- a/src/utils/other.py
+++ b/src/utils/other.py
@@ -74,7 +74,7 @@ def prepare_model_for_training(
        finetuning_type: str,
        output_embedding_layer_name: Optional[str] = "lm_head",
        use_gradient_checkpointing: Optional[bool] = True,
-        layer_norm_names: Optional[List[str]] = ["norm", "ln_f"] # for LLaMA and BLOOM setting
+        layer_norm_names: Optional[List[str]] = ["norm", "ln_f", "ln_attn", "ln_mlp"] # for LLaMA, BLOOM and Falcon settings
 ) -> PreTrainedModel:

    for name, param in model.named_parameters():