fix(llama): allow explicit head_dim when hidden_size not divisible by num_attention_heads
Author: Sriniketh24
Pull request: yes

## What this fixes

`LlamaConfig.validate_architecture()` unconditionally checks `hidden_size % num_attention_heads == 0`, but `LlamaAttention` already sizes its Q/K/V projections from `num_attention_heads * head_dim`, not from `hidden_size`. When `head_dim` is user-supplied the divisibility constraint is unnecessary and incorrectly blocks valid configs (issue #46082).

Reproduction (before fix):
```python
LlamaConfig(hidden_size=512, num_attention_heads=9, head_dim=56)
# StrictDataclassClassValidationError: The hidden size (512) is not a multiple of the number of attention heads (9).
```

## Approach

Track whether `head_dim` was user-supplied in `__post_init__` via a private `_head_dim_was_explicit` sentinel — set **before** the `None`-fallback so it captures caller intent rather than the derived value. Gate the `ValueError` on `not _head_dim_was_explicit`.

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

## Tests

Three new tests in `LlamaConfigTest`:
1. Explicit `head_dim` with non-divisible `hidden_size` → no error
2. No `head_dim`, non-divisible `hidden_size` → still raises
3. Save/reload roundtrip with explicit `head_dim` → works

```
python -m pytest tests/models/llama/test_modeling_llama.py::LlamaConfigTest -v
# 3 passed
```

Full suite (`tests/models/llama/test_modeling_llama.py`) shows 0 regressions vs `main` — 8 pre-existing FSDP/TP failures present identically on both branches.

Coordinated at: #46082 (no existing PR addresses this issue)

---

##### Was generative AI tooling used to co-author this PR?

- [X] Yes — Claude Code (claude-sonnet-4-6)

Generated-by: Claude Code (claude-sonnet-4-6) following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)

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

run-slow: llama

Hi @Sriniketh24 , I was glad to find this issue, as I am encountering the same error. I'm trying to open [kakaocorp/kanana-1.5-2.1b-instruct-2505](https://huggingface.co/kakaocorp/kanana-1.5-2.1b-instruct-2505) and face same error.

```py
model = AutoModelForCausalLM.from_pretrained(
    "kakaocorp/kanana-1.5-2.1b-instruct-2505",
    dtype="auto",
    device_map="auto",
)
```

```text
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
File .../.venv/lib/python3.12/site-packages/huggingface_hub/dataclasses.py:251, in strict.<locals>.wrap.<locals>.validate(self)
    250 try:
--> 251     validator(self)
    252 except (ValueError, TypeError) as e:

File .../.venv/lib/python3.12/site-packages/transformers/models/llama/configuration_llama.py:97, in LlamaConfig.validate_architecture(self)
     96 if self.hidden_size % self.num_attention_heads != 0:
---> 97     raise ValueError(
     98         f"The hidden size ({self.hidden_size}) is not a multiple of the number of attention "
     99         f"heads ({self.num_attention_heads})."
    100     )

ValueError: The hidden size (1792) is not a multiple of the number of attention heads (24).
```

Here is my understanding of the issue.

I checked the [`config.json`](https://huggingface.co/kakaocorp/kanana-1.5-2.1b-instruct-2505/blob/main/config.json) of [kakaocorp/kanana-1.5-2.1b-instruct-2505](https://huggingface.co/kakaocorp/kanana-1.5-2.1b-instruct-2505).

```json
{
    ...
    "head_dim": 128,
    "hidden_size": 1792,
    "num_attention_heads": 24,
    "num_key_value_heads": 8,
    "transformers_version": "4.48.3",
    ...
}
```

The configuration explicitly provides `head_dim=128`. Since the attention projections are sized using `num_attention_heads * head_dim`, `hidden_size` does not need to be divisible by `num_attention_heads` when `head_dim` is provided explicitly.

It seems that the current divisibility check should only apply when `head_dim` is omitted and derived from `hidden_size // num_attention_heads`.

I also checked when the current validation was introduced. It appears to be part of the strict configuration validation work added in 5.4.0. 

https://github.com/huggingface/transformers/blob/39f751a538ca67932cab53e6eb5763243674ae2c/src/transformers/models/llama/configuration_llama.py#L86-L100

So, rather than removing this validation, I think it would be better to align it with the actual constraints of `LlamaAttention`:

https://github.com/huggingface/transformers/blob/a9e2609950ac01ebc8b7a5e35b81c910e39d5a8b/src/transformers/models/llama/modeling_llama.py#L228-L249

https://github.com/huggingface/transformers/blob/a9e2609950ac01ebc8b7a5e35b81c910e39d5a8b/src/transformers/models/llama/modeling_llama.py#L209-L210

```py
def __post_init__(self, **kwargs):
    if self.head_dim is None:
        if self.hidden_size % self.num_attention_heads != 0:
            raise ValueError(
                f"The hidden size ({self.hidden_size}) is not a multiple of the number of attention "
                f"heads ({self.num_attention_heads})."
            )
        self.head_dim = self.hidden_size // self.num_attention_heads
    if self.num_key_value_heads is None:
        self.num_key_value_heads = self.num_attention_heads

    super().__post_init__(**kwargs)

def validate_architecture(self):
    """Part of `@strict`-powered validation. Validates the architecture of the config."""
    if num_attention_heads % num_key_value_heads != 0:
        raise ValueError("`num_attention_heads` must be divisible by `num_key_value_heads`.")
```

This would preserve the original goal of catching invalid configurations early while allowing valid configurations with an explicit `head_dim`.

And, as a minor aside, since current implementation supports all of MHA, GQA, and MQA... so maybe?

https://github.com/huggingface/transformers/blob/a9e2609950ac01ebc8b7a5e35b81c910e39d5a8b/src/transformers/models/llama/modeling_llama.py#L224-L226

@Sriniketh24 If you are not planning to continue with the fix, would you mind if I prepared a PR based on this approach?