[from_pretrained] imporve the error message when _no_split_modules …

…is not defined (huggingface#23861) * Better warning * Update src/transformers/modeling_utils.py Co-authored-by: Sylvain Gugger <[email protected]> * format line --------- Co-authored-by: Sylvain Gugger <[email protected]>
gojiteji · Jun 5, 2023 · 576755e · 576755e
1 parent 7e0aeec
commit 576755e
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
@@ -2734,7 +2734,10 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
                 target_dtype = torch.int8
 
             if model._no_split_modules is None:
-                raise ValueError(f"{model.__class__.__name__} does not support `device_map='{device_map}'` yet.")
+                raise ValueError(
+                    f"{model.__class__.__name__} does not support `device_map='{device_map}'`. To implement support, the model"
+                    "class needs to implement the `_no_split_modules` attribute."
+                )
             no_split_modules = model._no_split_modules
             if device_map not in ["auto", "balanced", "balanced_low_0", "sequential"]:
                 raise ValueError(