add upcast_lmhead option

Former-commit-id: 7ef69a1697c11ff13e7503360e40ef36cfb1c345
2024-01-19 23:54:25 +08:00
parent 1a39f529c0
commit be61bfda93
2 changed files with 5 additions and 1 deletions
--- a/src/llmtuner/hparams/model_args.py
+++ b/src/llmtuner/hparams/model_args.py
@@ -70,6 +70,10 @@ class ModelArguments:
        default=False,
        metadata={"help": "Whether or not to upcast the layernorm weights in fp32."}
    )
+    upcast_lmhead_output: Optional[bool] = field(
+        default=False,
+        metadata={"help": "Whether or not to upcast the output of lm_head in fp32."}
+    )
    hf_hub_token: Optional[str] = field(
        default=None,
        metadata={"help": "Auth token to log in with Hugging Face Hub."}