`_load_state_dict_into_meta_model` - `'NoneType' object has no attribute 'load_state_dict'`
Author: hlky
Pull request: no

https://github.com/huggingface/diffusers/actions/runs/13615360562/job/38057746315?pr=10898


```
model = StableDiffusionSafetyChecker(
  (vision_model): CLIPVisionModel(
    (vision_model): CLIPVisionTransformer(
      (emb...=1e-05, elementwise_affine=True)
    )
  )
  (visual_projection): Linear(in_features=32, out_features=64, bias=False)
)
state_dict = {'concept_embeds': tensor([[1., 1., 1.,  ..., 1., 1., 1.],
        [1., 1., 1.,  ..., 1., 1., 1.],
        [1., 1., 1...., 1., 1.,
         1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]]), 'special_care_embeds_weights': tensor([1., 1., 1.]), ...}
start_prefix = ''
expected_keys = ['concept_embeds', 'special_care_embeds', 'concept_embeds_weights', 'special_care_embeds_weights', 'vision_model.vision_model.embeddings.class_embedding', 'vision_model.vision_model.embeddings.patch_embedding.weight', ...]
device_map = None, offload_folder = None, offload_index = None
state_dict_folder = None, state_dict_index = None, dtype = torch.float16
hf_quantizer = None, is_safetensors = False, keep_in_fp32_modules = None
unexpected_keys = [], device_mesh = None
shard_file = '/github/home/.cache/huggingface/hub/models--hf-internal-testing--tiny-stable-diffusion-pipe/snapshots/3ee6c9f225f088ad5d35b624b6514b091e6a4849/safety_checker/pytorch_model.bin'

    @torch.no_grad()
    def _load_state_dict_into_meta_model(
        model: torch.nn.Module,
        state_dict: Dict[str, torch.Tensor],
        start_prefix,
        expected_keys,
        device_map=None,
        offload_folder=None,
        offload_index=None,
        state_dict_folder=None,
        state_dict_index=None,
        dtype=None,
        hf_quantizer=None,
        is_safetensors=False,
        keep_in_fp32_modules=None,
        unexpected_keys=None,  # passing `unexpected` for cleanup from quantization items
        device_mesh=None,
        shard_file=None,
    ):
        """
        This is somewhat similar to `_load_state_dict_into_model`, but deals with a model that has some or all of its
        params on a `meta` device. It replaces the model params with the data from the `state_dict`, while moving the
        params back to the normal device, but only for `loaded_state_dict_keys`.
    
        `start_prefix` is used for models which insert their name into model keys, e.g. `bert` in
        `bert.pooler.dense.weight`
    
        It also initialize tensor parallelism for each module if needed.
    
        """
        tensor_device = None
        if device_map is not None and device_map.get("", None) is not None:
            tensor_device = device_map[""].index if isinstance(device_map[""], torch.device) else device_map[""]
        if device_map is not None:
            device_map_regex = "|".join(sorted(device_map.keys(), reverse=True))
    
        # we need this later to initialize tensor parallelism
        if device_mesh is not None:
            full_tp_plan = model.config.base_model_tp_plan
            for submodule in model.modules():
                full_tp_plan.update(getattr(submodule, "_tp_plan", {}))
    
        file_pointer = None
        bin_state_dict = None
        if shard_file.endswith(".safetensors"):
            file_pointer = safe_open(shard_file, framework="pt", device=tensor_device)
        else:
            bin_state_dict = load_state_dict(shard_file, map_location="cpu")
    
        error_msgs = []
    
        is_quantized = hf_quantizer is not None
    
        is_torch_e4m3fn_available = hasattr(torch, "float8_e4m3fn")
    
        for serialized_param_name, empty_param in state_dict.items():
            # serialized_param_name is the raw, serialized name
            # fixed_param_name is the model's equivalent
            fixed_param_name, _ = model.rename_key(serialized_param_name)
    
            if fixed_param_name not in expected_keys:
                continue
    
            # we need to use serialized_param_name as file pointer is untouched
            param = (
                file_pointer.get_slice(serialized_param_name)
                if shard_file.endswith(".safetensors")
                else bin_state_dict[serialized_param_name]
            )
            # We convert floating dtypes to the `dtype` passed except for float8_e4m3fn type. We also want to keep the buffers/params
            # in int/uint/bool and not cast them.
            param_casting_dtype = None
            is_param_float8_e4m3fn = is_torch_e4m3fn_available and empty_param.dtype == torch.float8_e4m3fn
    
            if dtype is not None and empty_param.dtype.is_floating_point and not is_param_float8_e4m3fn:
                if (
                    keep_in_fp32_modules is not None
                    and keep_in_fp32_modules.search(fixed_param_name)
                    and dtype == torch.float16
                ):
                    param_casting_dtype = torch.float32
                else:
                    param_casting_dtype = dtype
    
            if device_mesh is not None:  # In this case, the param is already on the correct device!
                module_to_tp, param_type = find_submodule_and_param_name(model, fixed_param_name)
                current_module_plan = None
                full_tp_plan_ = "|".join(full_tp_plan.keys()).replace("*", "[0-9]+")
                if plan := re.search(full_tp_plan_, fixed_param_name):
                    match = re.sub("[0-9]+", "*", plan[0])
                    current_module_plan = full_tp_plan[match]
    
                if current_module_plan is not None:
                    tp_layer = translate_to_torch_parallel_style(current_module_plan)
                    rank = tensor_device
                    row, col = empty_param.shape
                    if "rowwise" == current_module_plan:
                        param = param[:, rank * (col // device_mesh.size()) : (rank + 1) * (col // device_mesh.size())]
                        shard = Shard(1)
                        tp_layer.desired_input_layouts = (Shard(-1),)
                    elif "colwise" == current_module_plan:
                        param = param[rank * (row // device_mesh.size()) : (rank + 1) * (row // device_mesh.size()), :]
                        shard = Shard(0)
                    else:
                        param = param[rank * (row // device_mesh.size()) : (rank + 1) * (row // device_mesh.size()), :]
                        shard = Shard(0)
                    if param_casting_dtype is not None and param_casting_dtype != empty_param.dtype:
                        param = param.to(param_casting_dtype)
                    local_parameter = DTensor.from_local(
                        param,
                        device_mesh=device_mesh,
                        placements=[shard] * device_mesh.ndim,
                    )
                    if isinstance(module_to_tp.weight, nn.Parameter):
                        local_parameter = torch.nn.Parameter(local_parameter)
                    module_to_tp.weight = local_parameter
                    input_fn = partial(tp_layer._prepare_input_fn, tp_layer.input_layouts, tp_layer.desired_input_layouts)
                    output_fn = partial(tp_layer._prepare_output_fn, tp_layer.output_layouts, tp_layer.use_local_output)
                    distribute_module(module_to_tp, device_mesh, None, input_fn, output_fn)
                else:
                    module_to_tp.load_state_dict({param_type: param[:]}, strict=False, assign=True)
    
            else:
                if device_map is None:
                    param_device = "cpu"
                else:
                    module_layer = re.search(device_map_regex, fixed_param_name)
                    if not module_layer:
                        raise ValueError(f"{fixed_param_name} doesn't have any device set.")
                    else:
                        param_device = device_map[module_layer.group()]
    
                if param_device == "disk":
                    if not is_safetensors:
                        offload_index = offload_weight(param[:], fixed_param_name, offload_folder, offload_index)
                elif param_device == "cpu" and state_dict_index is not None:
                    state_dict_index = offload_weight(param[:], fixed_param_name, state_dict_folder, state_dict_index)
                elif (
                    not is_quantized
                    or (not hf_quantizer.requires_parameters_quantization)
                    or (
                        not hf_quantizer.check_quantized_param(
                            model,
                            param,
                            fixed_param_name,
                            state_dict,
                            param_device=param_device,
                            device_map=device_map,
                        )
                    )
                ):
                    if is_fsdp_enabled():
                        param_device = "cpu" if is_local_dist_rank_0() else "meta"
                    module, param_type = find_submodule_and_param_name(model, fixed_param_name)
                    if param_casting_dtype is not None and param_casting_dtype != empty_param.dtype:
                        param = param[:].to(param_casting_dtype)
>                   module.load_state_dict(
                        {param_type: param[:].to(param_device)},
                        strict=False,
                        assign=True,
                    )
E                   AttributeError: 'NoneType' object has no attribute 'load_state_dict'
```

Comments:
Hi, this appears to be an error in `diffusers`. Are you sure `transformers` is the right repo?

Yes, the affected code is in `transformers` and the affected model is `StableDiffusionSafetyChecker` (in our case) which comes from `transformers`.

Here is the source of the problem
https://github.com/huggingface/transformers/blob/9e3a1072c25732d351a09218d8e8843148ecc3c9/src/transformers/modeling_utils.py#L707-L727

It is returning `None`.

It's related to a couple other issues that we picked up on in `diffusers`' tests.
https://github.com/huggingface/transformers/commit/a40f1ac602fe900281722254c52ce3773f28eb0e
https://github.com/huggingface/transformers/commit/dcbdf7e962c4b36140cc9ee76f870016121e69e5

cc @SunMarc 

Yes, I'm aware. This should fix the issue https://github.com/huggingface/transformers/pull/36512

Awesome, thanks for fast fix @SunMarc 

Closing now that #36512 is merged