fix kv cache

Former-commit-id: a9588e36e95bed896eea8d79ba7108447ff08f4b
2024-03-13 01:21:50 +08:00
parent 43fd80a1aa
commit 9a784fb4f3
2 changed files with 18 additions and 9 deletions
--- a/src/llmtuner/hparams/model_args.py
+++ b/src/llmtuner/hparams/model_args.py
@@ -101,6 +101,10 @@ class ModelArguments:
        default="offload",
        metadata={"help": "Path to offload model weights."},
    )
+    use_cache: bool = field(
+        default=True,
+        metadata={"help": "Whether or not to use KV cache in generation."},
+    )
    hf_hub_token: Optional[str] = field(
        default=None,
        metadata={"help": "Auth token to log in with Hugging Face Hub."},