transformers - 💡(How to fix) Fix LlamaConfig rejects explicit head_dim when hidden_size is not divisible by num_attention_heads [1 comments, 2 participants]

System Info

• transformers version: 5.8.0.dev0 / current main
• Checked against main commit a0fb01c6cda2301cb54de0efde5fac405836c4fe
• Python version: 3.13.12

Who can help?

@ArthurZucker @Cyrilvallez

Information

• The official example scripts
• My own modified scripts

Tasks

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

Reproduction

LlamaConfig exposes head_dim, but validation still rejects configs where hidden_size is not divisible by num_attention_heads, even when head_dim is explicitly provided.

from transformers import LlamaConfig, LlamaForCausalLM

config = LlamaConfig(
    vocab_size=99,
    hidden_size=512,
    intermediate_size=1024,
    num_hidden_layers=1,
    num_attention_heads=9,
    num_key_value_heads=1,
    head_dim=56,
)

model = LlamaForCausalLM(config)

This raises a validation error before the model can be constructed, even though the attention projection dimensions are well-defined from num_attention_heads * head_dim.

Expected behavior

If head_dim is explicitly provided, LlamaConfig should allow non-divisible hidden_size / num_attention_heads values and construct projections from num_attention_heads * head_dim.

If head_dim is not provided, the existing validation error should still be raised because head_dim must be derived from hidden_size // num_attention_heads.

Related custom-head_dim support has come up for other model families, for example #36659 and #37187.