fix #3238

Former-commit-id: 4d7e81ab4722d13bec6ca1af141f94bdc74d0883
2024-04-12 14:28:11 +08:00
parent 1a77de82fa
commit 31bbbb6d13
2 changed files with 6 additions and 3 deletions
--- a/src/llmtuner/hparams/parser.py
+++ b/src/llmtuner/hparams/parser.py
@@ -277,7 +277,11 @@ def get_infer_args(args: Optional[Dict[str, Any]] = None) -> _INFER_CLS:

    _verify_model_args(model_args, finetuning_args)

-    model_args.device_map = "auto"
+    if model_args.export_dir is not None:
+        model_args.device_map = {"": "cpu"}
+        model_args.compute_dtype = torch.float32
+    else:
+        model_args.device_map = "auto"

    return model_args, data_args, finetuning_args, generating_args