Fix DBRX MoE hidden size and expert GLU transposes

[kaixuanliu]

DbrxRouter/DbrxExperts are built from config.ffn_config, where hidden_size is the model hidden size and ffn_hidden_size the expert intermediate size, but the two were mixed up: the router linear used ffn_hidden_size inputs, DbrxExperts.forward reshaped hidden states with ffn_hidden_size, and the three matmuls in DbrxExpertGLU.forward had their transposes inverted. Any DBRX checkpoint fails in the MoE block on any device. This should be a regression introduced in #40132 (commit 7938e91) when the DBRX MoE modules were rewritten to read dims from config

DbrxModelTester also swapped the two sizes in ffn_config, which made the wrong lookups resolve to the right value and hid the bug; fixed as well.

@ArthurZucker @Cyrilvallez @vasqu pls help review, thx!

[kaixuanliu]

Here is an example to reproduce the bug:

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

tokenizer = AutoTokenizer.from_pretrained("Undi95/dbrx-base", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("Undi95/dbrx-base", device_map="auto", torch_dtype=torch.bfloat16)

input_text = "Databricks was founded in "
input_ids = tokenizer(input_text, return_tensors="pt").to(model.device)

outputs = model.generate(**input_ids, max_new_tokens=100)
print(tokenizer.decode(outputs[0]))

It will crash w/ log:

transformer.blocks.{0...39}.ffn.router.layer.weight | MISMATCH | Reinit due to size mismatch - ckpt: torch.Size([16, 6144]) vs model:torch.Size([16, 10752])

Notes:
- MISMATCH:     ckpt weights were loaded, but they did not match the original empty weight shapes.
Traceback (most recent call last):
  File "/root/upstream/transformers/dbrx.py", line 5, in <module>
    model = AutoModelForCausalLM.from_pretrained("Undi95/dbrx-base", device_map="auto", torch_dtype=torch.bfloat16)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/upstream/transformers/src/transformers/models/auto/auto_factory.py", line 402, in from_pretrained
    return model_class.from_pretrained(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/upstream/transformers/src/transformers/modeling_utils.py", line 4349, in from_pretrained
    loading_info = cls._finalize_model_loading(model, load_config, loading_info)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/upstream/transformers/src/transformers/modeling_utils.py", line 4525, in _finalize_model_loading
    log_state_dict_report(
  File "/root/upstream/transformers/src/transformers/utils/loading_report.py", line 278, in log_state_dict_report
    raise RuntimeError(
RuntimeError: You set `ignore_mismatched_sizes` to `False`, thus raising an error. For details look at the above report!

This PR can solve this bug.

[vasqu]

Approving on the assumption that we clarify my last comments

That is some wild mixup, checking slow ci in a second

[vasqu]

Do you know which PR introduced that mixup? That would be really nice to document

[kaixuanliu]

It should be introduced by #40132, and I updated the PR description.

[vasqu]

Not on you but personally prefer .T

[kaixuanliu]

updated.

[vasqu]

run-slow: dbrx

[github-actions]

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/dbrx"]
quantizations: []

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[github-actions]

CI Results

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	0d75e94b	workflow commit (merge commit)
PR	262a8bee	branch commit (from PR)
main	bd950935	base commit (on main)

✅ No failing test specific to this PR 🎉 👏 !

[vasqu]

Hmm test_tiny_model_logits still fails with resize issues, could you investigate? Might be a wrong format save because of random weights 🤔

[kaixuanliu]

Hmm test_tiny_model_logits still fails with resize issues, could you investigate? Might be a wrong format save because of random weights 🤔

Well this case also fails before this PR. In config.json ffn_config.moe_jitter_eps is the integer 0, we need to do some adjustment for type annotation in configuration_dbrx.py. And we need to make sure ffn_config.hidden_size must mirror the model's hidden_size, I have added related fix.

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: dbrx

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 30884851982:2
Result: success | Jobs: 8 | Tests: 376 | Failures: 0 | Duration: 2m 33s

[vasqu]

run-slow: dbrx

[github-actions]

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/dbrx"]
quantizations: []

[github-actions]

CI Results

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	f86fc4a0	workflow commit (merge commit)
PR	bc2e0ac6	branch commit (from PR)
main	f831d9c5	base commit (on main)

✅ No failing test specific to this PR 🎉 👏 !