Add custom head_dim support to Llama

[suhara]

What does this PR do?

Llama assumes that head_dim * num_heads == hidden_size and does not accommodate any models with custom head_dim size. This PR relaxes the assumption and makes Llama use custom head_dim sizes.

This PR has a dependency on the following PR:

• Use head_dim if in config for RoPE #32495
  • src/transformers/modeling_rope_utils.py to use config.head_dim for RoPE.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[amyeroberts]

Suggested change

	if config.head_dim is None:
	if (self.head_dim * self.num_heads) != self.hidden_size:
	raise ValueError(
	f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
	f" and `num_heads`: {self.num_heads})."
	)
	if config.head_dim is None and (self.head_dim * self.num_heads) != self.hidden_size:
	raise ValueError(
	f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
	f" and `num_heads`: {self.num_heads})."
	)

[suhara]

@amyeroberts Thanks for the suggestion! Updated the if block accordingly.

[amyeroberts]

cc @ArthurZucker

[ArthurZucker]

Hey! not sure we needs this (what I meant by a regression is that I thought we did allow head dim for llama) other models had this constraint lifted liike gemma I think

[suhara]

Hi @ArthurZucker

The motivation is that some custom Llama-architecture based models with custom head_dim sizes cannot be loaded by LlamaModel due to this constraint. (Some context: NVIDIA/NeMo#10078)

That said, I understand your concern on the regression issue. What would be your suggestion? If the existing Llama class is supposed to support official Llama models, creating a new class to cover custom Llama-based variant models would be an option?

[ArthurZucker]

Suggestions should fix the CI let's go with this.
Sorry for the delayed reviewed I was OOO for a bit

[ArthurZucker]

what you can do here is if head_dim is None: self.head_dim = self.hidden_size // self.num_heads

[suhara]

Thanks for the suggestion! Added.

[ArthurZucker]

Cool can you just run make fixup and make fix-copies

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

[suhara]

@ArthurZucker It seems that you created another PR and fix the remaining issues for this. Thank you!

[suhara]

Hi @ArthurZucker

I saw @Qubitium 's message in #32857. The newly created PR is missing the fix for o_proj.
I fixed the CI issues nd this PR should be ready to merge. Can you check?

Thanks!

[suhara]

Hi @ArthurZucker

The CI passed. Can you merge this PR (or #32857 after fixing the issue)? Thanks!

[suhara]

#32857 has been merged. Close this PR.