Fix config.softmax_scale not being considered - #3698
Conversation
| * ( | ||
| (1 / math.sqrt(self.q_head_dim)) | ||
| if self.config.softmax_scale is None | ||
| else self.config.softmax_scale |
There was a problem hiding this comment.
I'm not super familiar with the details of MLA... does it even make sense to have a configured softmax_scale for MLA? I'm wondering if we should instead fail in this case?
There was a problem hiding this comment.
We could ask the same about standard Attention, right? :)
One argument that objectively speaks for allowing this would be a hypothetical μP implementation for MLA, which would likely make use of softmax_scale similar to how it's done for standard Attention.
There was a problem hiding this comment.
For comparing this to standard attention, you can basically ignore the mscale stuff, that's just relevant for YaRN.
MLA defines its softmax_scale as (qk_head_dim + qk_pos_emb_head_dim)^-0.5. The first term in the sum is the same as in standard attention; the second term is the per-head dimensionality of the decoupled queries, which are part of the modified, decoupled RoPE for MLA.
900a73f to
1922754
Compare
|
Fixed other places where it was missing and added non-functional unit tests. |
1922754 to
e328f83
Compare
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
5 participants
Previously,
config.softmax_scalewas sometimes ignored and the default used unconditionally. This is fine in most cases, but whensoftmax_scaleis set, we do – of course – want to use it.