423A35C7/LLaMA-Factory
1
0
Fork 0
mirror of https://github.com/hiyouga/LLaMA-Factory.git synced 2025-08-06 05:32:50 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix #2127

Browse Source 
Former-commit-id: ebee4f6a2a20a39ad7fbb87edfa63e244da7b5e6
This commit is contained in:
hiyouga 2024-01-09 14:49:13 +08:00
parent de9148930f
commit f53bc7d9a0
1 changed files with 2 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
src/llmtuner/model/utils.py
View File

@ -27,14 +27,12 @@ def dispatch_model(model: "PreTrainedModel") -> "PreTrainedModel":
if ( if (
torch.cuda.device_count() > 1 torch.cuda.device_count() > 1
and isinstance(model, PreTrainedModel) and isinstance(model, PreTrainedModel)
and getattr(model.config, "model_type", None) != "chatglm" and model._no_split_modules is not None
and model.config.model_type != "chatglm"
): ):
from accelerate import dispatch_model from accelerate import dispatch_model
from accelerate.utils import infer_auto_device_map, get_balanced_memory from accelerate.utils import infer_auto_device_map, get_balanced_memory
if getattr(model, "_no_split_modules", None) is None:
raise ValueError("The model class needs to implement the `_no_split_modules` attribute.")
kwargs = {"dtype": model.dtype, "no_split_module_classes": model._get_no_split_modules("auto")} kwargs = {"dtype": model.dtype, "no_split_module_classes": model._get_no_split_modules("auto")}
max_memory = get_balanced_memory(model, **kwargs) max_memory = get_balanced_memory(model, **kwargs)
# Make sure tied weights are tied before creating the device map. # Make sure tied weights are tied before creating the device map.