use existing is_local_dist_rank_0 fn to check devices

src/transformers/modeling_utils.py

@@ -933,7 +933,7 @@ def _load_state_dict_into_meta_model(
             )
         ):
             if is_fsdp_enabled():
-                param_device = "cpu" if device_map[""] == 0 else "meta"
+                param_device = "cpu" if is_local_dist_rank_0() else "meta"
             # For backward compatibility with older versions of `accelerate` and for non-quantized params
             set_module_tensor_to_device(model, param_name, param_device, **set_module_kwargs)
         else:
@@ -945,7 +945,7 @@ def _load_state_dict_into_meta_model(
                 module, tensor_name = get_module_from_name(model, param_name)
                 value = getattr(module, tensor_name)
                 param_to = "cpu"
-                if is_fsdp_enabled() and value.data.device.index != 0:
+                if is_fsdp_enabled() and not is_local_dist_rank_0():
                     param_to = "meta"
                 value = type(value)(value.data.to(param_to), **value.__dict__)
                 setattr(module, tensor_name, value)