from_pretrained is reporting Missking Keys even though they are handled in load_state_dict_pre_hook
Author: ZhiyuanChen
Pull request: no
Labels: name: Usage

name: bug

name: Mixture of Experts

### System Info

- `transformers` version: 4.44.2
- Platform: macOS-15.1-arm64-arm-64bit
- Python version: 3.10.14
- Huggingface_hub version: 0.23.3
- Safetensors version: 0.4.3
- Accelerate version: 0.34.2
- Accelerate config:    not found
- PyTorch version (GPU?): 2.4.1 (False)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

```python
class XxxSparseLayer(nn.Module):
    def __init__(self, config: XxxConfig):
        super().__init__()
        self.num_experts = config.num_experts
        self.top_k = config.expert_top_k
        self.attention = XxxAttention(config)
        self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.router = nn.Linear(config.hidden_size, self.num_experts)
        self.experts = nn.ModuleList([XxxFeedForward(config) for _ in range(self.num_experts)])
        self._register_load_state_dict_pre_hook(self.copy_ffn_params)

    def forward(
        self,
        hidden_states: Tensor,
        attention_mask: torch.FloatTensor | None = None,
        head_mask: torch.FloatTensor | None = None,
        output_attentions: bool = False,
    ) -> Tuple[Tensor, ...]:
        self_attention_outputs = self.attention(
            hidden_states,
            attention_mask,
            head_mask,
            output_attentions=output_attentions,
        )
        attention_output = self_attention_outputs[0]
        layer_output = self.layer_norm(attention_output)

        router_logits = self.router(layer_output[:, 0, :])
        router_probs = router_logits.softmax(dim=-1)
        router_weights, router_idx = router_probs.topk(self.top_k, dim=-1)
        router_weights /= router_weights.sum(dim=-1, keepdim=True)

        expert_outputs = torch.stack([self.experts[i](layer_output) for i in range(self.num_experts)], dim=1)
        solicited_outputs = expert_outputs[torch.arange(router_idx.size(0)).unsqueeze(1), router_idx]
        weighted_outputs = (solicited_outputs * router_weights.unsqueeze(-1).unsqueeze(-1)).sum(1)

        layer_output = weighted_outputs + layer_output
        return (layer_output,) + self_attention_outputs[1:]

    def copy_ffn_params(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
        ffn_prefix = prefix + "ffn."
        ffn_states = {k: v for k, v in state_dict.items() if k.startswith(ffn_prefix)}
        ffn_states = {k: v for k, v in state_dict.items() if "layer_norm" not in k}
        for k, v in ffn_states.items():
            for i in range(self.num_experts):
                state_dict[k.replace("ffn.", f"experts.{i}.")] = v.clone()
            del state_dict[k]
```

### Expected behavior

I'm trying to upcycle a dense model to a sparsely activated MoE. 

The common practice is to copy the weights of FFN for each sub-expert FFN. 

PyTorch has implemented a `_register_load_state_dict_pre_hook` so that the renaming and copying process can happen at loading. 

However, while the weights are successfully loaded (I confirmed it through manual inspection), transformers still complains the following (they would raise an issue if using torch.load).

```
Some weights of XxxForSequencePrediction were not initialized from the model checkpoint at xxx and are newly initialized: ['xxx.encoder.layer.0.experts.0.in_proj.bias', 'xxx.encoder.layer.0.experts.0.in_proj.weight', 'xxx.encoder.layer.0.experts.0.out_proj.bias', 'xxx.encoder.layer.0.experts.0.out_proj.weight', 'xxx.encoder.layer.0.experts.1.in_proj.bias', 'xxx.encoder.layer.0.experts.1.in_proj.weight', 'xxx.encoder.layer.0.experts.1.out_proj.bias', 'xxx.encoder.layer.0.experts.1.out_proj.weight', 'xxx.encoder.layer.0.router.bias', 'xxx.encoder.layer.0.router.weight', 'xxx.encoder.layer.1.experts.0.in_proj.bias', 'xxx.encoder.layer.1.experts.0.in_proj.weight', 'xxx.encoder.layer.1.experts.0.out_proj.bias', 'xxx.encoder.layer.1.experts.0.out_proj.weight', 'xxx.encoder.layer.1.experts.1.in_proj.bias', 'xxx.encoder.layer.1.experts.1.in_proj.weight', 'xxx.encoder.layer.1.experts.1.out_proj.bias', 'xxx.encoder.layer.1.experts.1.out_proj.weight', 'xxx.encoder.layer.1.router.bias', 'xxx.encoder.layer.1.router.weight', 'xxx.encoder.layer.10.experts.0.in_proj.bias', 'xxx.encoder.layer.10.experts.0.in_proj.weight', 'xxx.encoder.layer.10.experts.0.out_proj.bias', 'xxx.encoder.layer.10.experts.0.out_proj.weight', 'xxx.encoder.layer.10.experts.1.in_proj.bias', 'xxx.encoder.layer.10.experts.1.in_proj.weight', 'xxx.encoder.layer.10.experts.1.out_proj.bias', 'xxx.encoder.layer.10.experts.1.out_proj.weight', 'xxx.encoder.layer.10.router.bias', 'xxx.encoder.layer.10.router.weight', 'xxx.encoder.layer.11.experts.0.in_proj.bias', 'xxx.encoder.layer.11.experts.0.in_proj.weight', 'xxx.encoder.layer.11.experts.0.out_proj.bias', 'xxx.encoder.layer.11.experts.0.out_proj.weight', 'xxx.encoder.layer.11.experts.1.in_proj.bias', 'xxx.encoder.layer.11.experts.1.in_proj.weight', 'xxx.encoder.layer.11.experts.1.out_proj.bias', 'xxx.encoder.layer.11.experts.1.out_proj.weight', 'xxx.encoder.layer.11.router.bias', 'xxx.encoder.layer.11.router.weight', 'xxx.encoder.layer.2.experts.0.in_proj.bias', 'xxx.encoder.layer.2.experts.0.in_proj.weight', 'xxx.encoder.layer.2.experts.0.out_proj.bias', 'xxx.encoder.layer.2.experts.0.out_proj.weight', 'xxx.encoder.layer.2.experts.1.in_proj.bias', 'xxx.encoder.layer.2.experts.1.in_proj.weight', 'xxx.encoder.layer.2.experts.1.out_proj.bias', 'xxx.encoder.layer.2.experts.1.out_proj.weight', 'xxx.encoder.layer.2.router.bias', 'xxx.encoder.layer.2.router.weight', 'xxx.encoder.layer.3.experts.0.in_proj.bias', 'xxx.encoder.layer.3.experts.0.in_proj.weight', 'xxx.encoder.layer.3.experts.0.out_proj.bias', 'xxx.encoder.layer.3.experts.0.out_proj.weight', 'xxx.encoder.layer.3.experts.1.in_proj.bias', 'xxx.encoder.layer.3.experts.1.in_proj.weight', 'xxx.encoder.layer.3.experts.1.out_proj.bias', 'xxx.encoder.layer.3.experts.1.out_proj.weight', 'xxx.encoder.layer.3.router.bias', 'xxx.encoder.layer.3.router.weight', 'xxx.encoder.layer.4.experts.0.in_proj.bias', 'xxx.encoder.layer.4.experts.0.in_proj.weight', 'xxx.encoder.layer.4.experts.0.out_proj.bias', 'xxx.encoder.layer.4.experts.0.out_proj.weight', 'xxx.encoder.layer.4.experts.1.in_proj.bias', 'xxx.encoder.layer.4.experts.1.in_proj.weight', 'xxx.encoder.layer.4.experts.1.out_proj.bias', 'xxx.encoder.layer.4.experts.1.out_proj.weight', 'xxx.encoder.layer.4.router.bias', 'xxx.encoder.layer.4.router.weight', 'xxx.encoder.layer.5.experts.0.in_proj.bias', 'xxx.encoder.layer.5.experts.0.in_proj.weight', 'xxx.encoder.layer.5.experts.0.out_proj.bias', 'xxx.encoder.layer.5.experts.0.out_proj.weight', 'xxx.encoder.layer.5.experts.1.in_proj.bias', 'xxx.encoder.layer.5.experts.1.in_proj.weight', 'xxx.encoder.layer.5.experts.1.out_proj.bias', 'xxx.encoder.layer.5.experts.1.out_proj.weight', 'xxx.encoder.layer.5.router.bias', 'xxx.encoder.layer.5.router.weight', 'xxx.encoder.layer.6.experts.0.in_proj.bias', 'xxx.encoder.layer.6.experts.0.in_proj.weight', 'xxx.encoder.layer.6.experts.0.out_proj.bias', 'xxx.encoder.layer.6.experts.0.out_proj.weight', 'xxx.encoder.layer.6.experts.1.in_proj.bias', 'xxx.encoder.layer.6.experts.1.in_proj.weight', 'xxx.encoder.layer.6.experts.1.out_proj.bias', 'xxx.encoder.layer.6.experts.1.out_proj.weight', 'xxx.encoder.layer.6.router.bias', 'xxx.encoder.layer.6.router.weight', 'xxx.encoder.layer.7.experts.0.in_proj.bias', 'xxx.encoder.layer.7.experts.0.in_proj.weight', 'xxx.encoder.layer.7.experts.0.out_proj.bias', 'xxx.encoder.layer.7.experts.0.out_proj.weight', 'xxx.encoder.layer.7.experts.1.in_proj.bias', 'xxx.encoder.layer.7.experts.1.in_proj.weight', 'xxx.encoder.layer.7.experts.1.out_proj.bias', 'xxx.encoder.layer.7.experts.1.out_proj.weight', 'xxx.encoder.layer.7.router.bias', 'xxx.encoder.layer.7.router.weight', 'xxx.encoder.layer.8.experts.0.in_proj.bias', 'xxx.encoder.layer.8.experts.0.in_proj.weight', 'xxx.encoder.layer.8.experts.0.out_proj.bias', 'xxx.encoder.layer.8.experts.0.out_proj.weight', 'xxx.encoder.layer.8.experts.1.in_proj.bias', 'xxx.encoder.layer.8.experts.1.in_proj.weight', 'xxx.encoder.layer.8.experts.1.out_proj.bias', 'xxx.encoder.layer.8.experts.1.out_proj.weight', 'xxx.encoder.layer.8.router.bias', 'xxx.encoder.layer.8.router.weight', 'xxx.encoder.layer.9.experts.0.in_proj.bias', 'xxx.encoder.layer.9.experts.0.in_proj.weight', 'xxx.encoder.layer.9.experts.0.out_proj.bias', 'xxx.encoder.layer.9.experts.0.out_proj.weight', 'xxx.encoder.layer.9.experts.1.in_proj.bias', 'xxx.encoder.layer.9.experts.1.in_proj.weight', 'xxx.encoder.layer.9.experts.1.out_proj.bias', 'xxx.encoder.layer.9.experts.1.out_proj.weight', 'xxx.encoder.layer.9.router.bias', 'xxx.encoder.layer.9.router.weight', 'xxx.pooler.dense.bias', 'xxx.pooler.dense.weight', 'sequence_head.decoder.bias', 'sequence_head.decoder.weight']
```

Comments:
Thanks for the issue @ZhiyuanChen, cc @ArthurZucker for MoEs

[#33680 ](https://github.com/huggingface/transformers/issues/33680)
This looks similiar to this where torch.load is the problem

Would try to troubleshoot . Can you pls share your insights so it would be helpful

> [#33680 ](https://github.com/huggingface/transformers/issues/33680) This looks similiar to this where torch.load is the problem
> 
> Would try to troubleshoot . Can you pls share your insights so it would be helpful

I tried to load with `torch.load`, and pytorch is fine (though with different error messages, that's why I created #33680)


I think I'll be able to submit a patch for this next week, as I'm super busy with ICLR deadline... 
But if anyone can't wait for it, I think this is resulted by the way `_IncompatibleKeys` are handled. 

Specifically, PyTorch implemented a `_register_load_state_dict_pre_hook` and a `register_load_state_dict_post_hook` to allow you modify the `state_dict` or `incompatible_keys ` on the fly (that's why they are mutual objects so you can modify them inlace). (Note that the leading `_` in pre-hook does not mean it's a private API, it should mean it is called on a module level, well the post-hook is only called once on the top-model level.) 

As a result, you cannot trust the key mismatch by just inspecting them externally. The only reliable way to get the correct `_IncompatibleKeys` is to get them from the `load_state_dict` call. (Technically speaking, you can manually call these hooks, but as they are inplace ops, remember to make a deep copy of the state dict first). 

Currently, transformers seems to inspect the `incompatible_keys` externally, so both hooks are ignored. 

I'll see if i can solve this issue

Thanks for the insights @ZhiyuanChen 

We'd gladly welcome a PR to help fix this @ZhiyuanChen!

@ZhiyuanChen  May i get the XxxConfig you used please, i am trying to debug this
Thanks

> @ZhiyuanChen May i get the XxxConfig you used please, i am trying to debug this Thanks

Sorry, it is a custom model under review. 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.