support export size setting

Former-commit-id: 1a4de54586c21cdbbc89f8a716ca5a54c87a6120
2023-11-26 18:34:09 +08:00
parent 953a562ec1
commit 569860d7ac
4 changed files with 11 additions and 5 deletions
--- a/src/llmtuner/hparams/finetuning_args.py
+++ b/src/llmtuner/hparams/finetuning_args.py
@@ -149,6 +149,10 @@ class FinetuningArguments(FreezeArguments, LoraArguments, RLHFArguments):
        default=None,
        metadata={"help": "Path to the directory to save the exported model."}
    )
+    export_size: Optional[int] = field(
+        default=1,
+        metadata={"help": "The file shard size (in GB) of the exported model."}
+    )
    plot_loss: Optional[bool] = field(
        default=False,
        metadata={"help": "Whether to plot the training loss after fine-tuning or not."}