Fix loading with only 1 device or distributed config

src/transformers/core_model_loading.py

@@ -38,12 +38,12 @@
 
 
 _torch_distributed_available = torch.distributed.is_available()
+if _torch_distributed_available:
+    from torch.distributed.tensor import DTensor
 
 if TYPE_CHECKING:
     from .modeling_utils import LoadStateDictConfig, PreTrainedModel
     from .quantizers import HfQuantizer
-elif _torch_distributed_available:
-    from torch.distributed.tensor import DTensor
 
 logger = get_logger(__name__)
 
@@ -799,7 +799,9 @@ def materialize_tensors(self) -> dict[str, list[torch.Tensor]]:
                 tensors = [future.result() for future in tensors if future.result() is not None]
             # Sync loading
             elif callable(tensors[0]):
-                tensors = [tensor for func in tensors if (tensor := func()) is not None]
+                tensors = [func() for func in tensors]
+                # Some may be None for some distributed setups
+                tensors = [tensor for tensor in tensors if tensor is not None]
             # Add them to the new dictionary
             collected_tensors[key] = tensors
 
@@ -1019,9 +1021,10 @@ def _job():
 
     if thread_pool is not None:
         return thread_pool.submit(_job)
-    # Return the Callable here, not the Tensor itself, so we actually delay loading
-    # to avoid saturating cpu memory during Conversion
-    return _job
+    else:
+        # Return the Callable here, not the Tensor itself, so we actually delay loading to avoid saturating cpu
+        # memory during Conversion
+        return _job
 
 
 def dot_natural_key(s: str):
@@ -1307,21 +1310,9 @@ def convert_and_load_state_dict_in_model(
     """
     base_model_prefix = model.base_model_prefix
     tp_plan = tp_plan or {}
+    device_map = load_config.device_map or {"": "cpu"}
     hf_quantizer = load_config.hf_quantizer
     dtype = load_config.dtype
-    device_mesh = load_config.device_mesh
-
-    if load_config.device_map is not None:
-        device_map = load_config.device_map
-    elif device_mesh is not None:
-        if device_mesh.device_type == "cpu":
-            device_map = {"": torch.device("cpu")}
-        else:
-            device_map = {
-                "": torch.device(device_mesh.device_type, getattr(torch, device_mesh.device_type).current_device())
-            }
-    else:
-        device_map = {"": "cpu"}
     disk_offload_folder = load_config.disk_offload_folder
     offload_buffers = load_config.offload_buffers
     dtype_plan = load_config.dtype_plan or {}

src/transformers/modeling_utils.py

@@ -4188,6 +4188,12 @@ def from_pretrained(
 
         if distributed_config is not None:
             device_mesh = init_device_mesh(distributed_config)
+            # When using any distributed setup, we simply set the device_map here to the current rank so that we correctly
+            # load the params on the right device
+            if device_mesh.device_type == "cpu":
+                device_map = {"": torch.device("cpu")}
+            else:
+                device_map = {"": torch.device(device_mesh.device_type, int(os.environ["LOCAL_RANK"]))}
         else:
             # Accelerate path
             if device_map == "auto" and int(os.environ.get("WORLD_SIZE", "0")):
@@ -4342,10 +4348,9 @@ def from_pretrained(
 
         if distributed_config is not None:
             model = distribute_model(model, distributed_config, device_mesh)
-        else:
-            # Accelerate path: auto device mapping
-            if device_map is not None:
-                device_map = _get_device_map(model, device_map, max_memory, hf_quantizer)
+        elif device_map is not None:
+            # Expand device_map if it was passed as a `str`, i.e. `device_map="auto"`
+            device_map = _get_device_map(model, device_map, max_memory, hf_quantizer)
 
         # Finalize model weight initialization
         active_tp_plan = (
@@ -4907,7 +4912,7 @@ def mark_tied_weights_as_initialized(self, loading_info):
         later as they will be tied (overwritten) anyway.
         This is very important as most embeddings are tied, and they are huge params (vocabularies are often 256k), so
         running inits on them is very costly."""
-        for tied_param in self.all_tied_weights_keys.keys():
+        for tied_param in getattr(self, "all_tied_weights_keys", {}).keys():
             param = self.get_parameter(tied_param)
             setattr(param, "_is_hf_initialized", True)
 

tests/utils/test_modeling_utils.py

@@ -442,6 +442,7 @@ def test_model_from_pretrained_fsdp_distributes_before_loading(self):
             fake_mesh = mock.Mock()
             fake_mesh.mesh_dim_names = ("fsdp",)
             fake_mesh.ndim = 1
+            fake_mesh.device_type = "cpu"
 
             def fake_apply_fsdp(model, fsdp_mesh, fsdp_plan):
                 call_order.append("distribute")