fix export format

Former-commit-id: 7c82bd396b9e6ff395850ad544d95cbf1b7557cd
2023-12-28 18:40:46 +08:00
parent 53d7c5109f
commit af3f5b6e16
3 changed files with 29 additions and 1 deletions
--- a/src/llmtuner/train/tuner.py
+++ b/src/llmtuner/train/tuner.py
@@ -47,7 +47,11 @@ def export_model(args: Optional[Dict[str, Any]] = None):

    model.config.use_cache = True
    model = model.to("cpu")
-    model.save_pretrained(model_args.export_dir, max_shard_size="{}GB".format(model_args.export_size))
+    model.save_pretrained(
+        save_directory=model_args.export_dir,
+        max_shard_size="{}GB".format(model_args.export_size),
+        safe_serialization=(not model_args.export_lecacy_format)
+    )

    try:
        tokenizer.padding_side = "left" # restore padding side