From 19d33ede137a687417bb8697f10e1749781d0d23 Mon Sep 17 00:00:00 2001 From: hiyouga Date: Sun, 4 Feb 2024 15:51:47 +0800 Subject: [PATCH] fix #2420 --- src/llmtuner/model/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/llmtuner/model/utils.py b/src/llmtuner/model/utils.py index 7ff7aa34..9d45c290 100644 --- a/src/llmtuner/model/utils.py +++ b/src/llmtuner/model/utils.py @@ -41,7 +41,7 @@ def dispatch_model(model: "PreTrainedModel") -> "PreTrainedModel": # Make sure tied weights are tied before creating the device map. model.tie_weights() device_map = infer_auto_device_map(model, max_memory=max_memory, **kwargs) - device_map_kwargs = {"device_map": device_map} + device_map_kwargs = {"device_map": device_map, "offload_dir": "offload"} if "skip_keys" in inspect.signature(dispatch_model).parameters: device_map_kwargs["skip_keys"] = model._skip_keys_device_placement return dispatch_model(model, **device_map_kwargs)