align shape suffix between attn and linear attn - #4390
Conversation
meta-cla
Bot
added
the
CLA Signed
|
|
||
| # Shape suffixes: | ||
| # B = singleton kernel batch, T = packed tokens, H = attention heads, | ||
| # D = projection rows per head, I = input features. |
There was a problem hiding this comment.
what is I, compared with D and F, in attention.py, and in
https://github.com/pytorch/torchtitan/blob/main/torchtitan/models/common/moe.py#L27
There was a problem hiding this comment.
D and F refers to model dim and FFN hidden dim. I is input feature dim/input width.
Here in this file
There was a problem hiding this comment.
couldn't understand, you have two I's in side the same model, how could it be used as identifier?
There was a problem hiding this comment.
_scale_mla_heads
q_projection and kv_projection twice, whose input dim are different.
| scales_N11 = _replicated_scales(scales_N, weight).view(-1, 1, 1) | ||
| heads_NDI = weight.view(num_heads, rows_per_head, weight.shape[1]) | ||
| heads_NDI[:, :nope_rows_per_head].mul_(scales_N11.pow(nope_scale_exponent)) | ||
| scales_H11 = _replicated_scales(scales_H, weight).view(-1, 1, 1) |
There was a problem hiding this comment.
why 1 here, but B = singleton kernel batch
There was a problem hiding this comment.
the 1 here is not batch dim. It's from .view(-1, 1, 1), to multiply with heads_HDI.
There was a problem hiding this comment.
what is "singleton kernel batch" and is it not always 1?
There was a problem hiding this comment.
B is always 1 here since we removed batch dim, it is needed only because kda still has the batch dim. #4164 (comment)
There was a problem hiding this comment.
@drisspg
I saw KDA kernels accepting inputs with both batch (B) and tokens (T). Curious that if the inputs are restricted to only have T dim, would it be further optimized?
2 participants
PR title
Align attention shape suffixes across TorchTitan
Summary
Stack created with GitHub Stacks CLI • Give Feedback 💬