improve aqlm optim

Former-commit-id: 81be999b407e988c2f42764d827ac859d079ed3e
2024-03-05 20:49:50 +08:00
parent a10bead9b5
commit 46ee267cfc
4 changed files with 7 additions and 3 deletions
--- a/src/llmtuner/hparams/model_args.py
+++ b/src/llmtuner/hparams/model_args.py
@@ -121,6 +121,9 @@ class ModelArguments:
        default=False,
        metadata={"help": "For debugging purposes, print the status of the parameters in the model."},
    )
+    aqlm_optimization: Optional[bool] = field(
+        default=False, metadata={"help": "Whether or not to optimize the training performance of AQLM models."}
+    )

    def __post_init__(self):
        self.compute_dtype = None