LlamaConfig rejects explicit head_dim when hidden_size is not divisible by num_attention_heads
Author: Anri-Lombard
Pull request: no

### 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
- [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

`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.

```python
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.


Comments:
Dug into this a bit.

The modeling code in `LlamaAttention` is already fine, projections are sized from `num_attention_heads * head_dim`, not `hidden_size`. The issue is in `configuration_llama.py`.

`validate_architecture` fires the divisibility check unconditionally, but by the time it runs (via `@strict`, post `__init__`), `head_dim` is always set, either from the user or derived as a fallback. So you can't just check `head_dim is not None` in the validator.

Fix is to capture intent early in `__post_init__`:

```python
def __post_init__(self):
    self._head_dim_was_explicit = self.head_dim is not None
    if self.head_dim is None:
        self.head_dim = self.hidden_size // self.num_attention_heads
```

Then gate the check on it:

```python
if self.hidden_size % self.num_attention_heads != 0 and not self._head_dim_was_explicit:
    raise ValueError(...)
```

One thing worth testing: save/reload the config. `_head_dim_was_explicit` won't be in the JSON, but `head_dim` will, so `__post_init__` should set the flag correctly on reload anyway. Happy to open a PR if that's useful.

I've implemented a fix on my fork and would like to open a PR.

**Approach** — same direction as @matdou's analysis: track caller intent in  before the  fallback:

```python
def __post_init__(self, **kwargs):
    self._head_dim_was_explicit = self.head_dim is not None  # capture intent before fallback
    if self.head_dim is None:
        self.head_dim = self.hidden_size // self.num_attention_heads
    ...

def validate_architecture(self):
    if self.hidden_size % self.num_attention_heads != 0 and not self._head_dim_was_explicit:
        raise ValueError(...)
```

The sentinel is a private attribute (not a dataclass field), so it's excluded from JSON serialization. On save/reload, `head_dim` is present in the JSON so `__post_init__` correctly sets `_head_dim_was_explicit = True` again — no round-trip regression.

**Tests verified locally:**
- Explicit `head_dim=56` with `hidden_size=512, num_heads=9` → no error ✓
- No `head_dim`, non-divisible `hidden_size` → still raises ✓  
- Save/reload roundtrip with explicit `head_dim` → works ✓
- Full `tests/models/llama/test_modeling_llama.py` — 0 regressions (8 pre-existing FSDP/TP failures unrelated to this change) ✓

Happy to open the PR if this approach looks good. No existing PR addresses this issue.

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.