Minor llama4 fixes

[molbap]

What does this PR do?

Fixes a couple nits/wrong defaults in Llama4 code for DynamicCache init - not sure we should even default to it, I think.
The scaling however has to be passed else `eager_attention_forward" will fail.

[github-actions]

Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. The CI will be paused while the PR is in draft mode. When it is ready for review, please click the Ready for review button (at the bottom of the PR page). This will assign reviewers and trigger CI.

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

[ArthurZucker]

super weird that this still work, can you run the test / check where this is from? (blame) might be expected 😐

[molbap]

I think it works with sdpa and other kernels because they default the scaling to sqrt(head_dim) when None is passed, checking

[molbap]

changing tests, they were broken too 😀

[molbap]

I ran the tests, it seems it works. I think it was broken during the interim after the recent VLMs refactor

[molbap]

...but, because of TP, setting the scaling manually here is probably something that could break the vision attention module. I've set a vision_eager_attention_forward to at least fix the eager forward that is currently broken, while preserving the behaviour we had and err on the side of caution. @ArthurZucker alright with you?

[vasqu]

Oh man, I see that we use the config's head dim and it's only set and done so in the text attention 👀

So in essence, the local head_dim was previously used in the vision attention. But depending on the TP degree, this is different. Did I understand it correctly? Imo, this is still a mistake, and not worth being BC: We essentially have different models based on different TP degrees - seems unreasonable to me and not the point of TP.

[vasqu]

Hmm, looking into the code

transformers/src/transformers/models/llama4/modeling_llama4.py

Lines 1065 to 1074 in 0f77ca7

	class Llama4VisionAttention(nn.Module):
	def __init__(self, config: Llama4VisionConfig):
	super().__init__()
	self.config = config
	self.embed_dim = config.hidden_size
	self.num_heads = config.num_attention_heads
	self.head_dim = config.hidden_size // config.num_attention_heads
	self.num_key_value_groups = 1
	self.attention_dropout = config.attention_dropout
	self.scaling = self.head_dim**-0.5

It seems to me, the global head_dim was previously used (by module.head_dim) before #37576 - not sure what the local head_dim would be.

[molbap]

I think None was passed (self.scaling was unset before the PR you linked). Definitely agree lol, the whole point of tp_plan is to have one model for all possible TP degrees. Passing scaling=None should be safe in all cases here, what was broken was mainly eager_attention_forward that did not compute scaling but instead used only the passed one, as you say... so now in main if you use eager you get a float * NoneType multiplication error here.

So sdpa/flex/etc are fine IMO, just the eager path is currently broken, one simple way that should be DTensor-compatible when it comes to scaling could be e.g. scaling = 1.0 / torch.sqrt(torch.tensor(query.size(-1), dtype=query.dtype, device=query.device)) WDYT?

[vasqu]

Shouldn't we be able to just pass self.scaling here @

transformers/src/transformers/models/llama4/modeling_llama4.py

Line 1120 in 34c1e29

scaling=None,

But yea, if that doesn't work, then I'm definitely pro your solution :D

[molbap]

yeah that was my initial fix here 9ffde0b! but I wasn't sure of why things were in this state, it's good to put in writing 🤗

[ArthurZucker]

LGTM thanks for digging!
@vasqu happy if you upstream the changes to our interface to better support scaling / default etc to make it more user friendly and avoid us these kind of issues!

[vasqu]

Sounds good :D