You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Currently, the TinyLlama model fails with fused layers because the shapes are not set up correctly for that specific model due to not handling n_kv_heads.
The general calculation seems to be: n_tokens*(n_heads+(n_kv_heads*2))*(hidden_size // n_heads)
This should fix problems for all models with GQA.
EDIT: Potentially also fix shapes during forward of GEMM.
print(x.reshape(-1, x.shape[-1]).shape[1] / self.scales.shape[0]) is 4 for TinyLlama, needs to be multiple of 32
EDIT: Probably best to take inspiration from Llama / CodeLlama
Currently, the TinyLlama model fails with fused layers because the shapes are not set up correctly for that specific model due to not handling
n_kv_heads
.The general calculation seems to be:
n_tokens*(n_heads+(n_kv_heads*2))*(hidden_size // n_heads)
This should fix problems for all models with GQA.
EDIT: Potentially also fix shapes during forward of GEMM.
print(x.reshape(-1, x.shape[-1]).shape[1] / self.scales.shape[0])
is 4 for TinyLlama, needs to be multiple of 32EDIT: Probably best to take inspiration from Llama / CodeLlama
No kv_heads:
https://github.com/facebookresearch/llama/blob/1076b9c51c77ad06e9d7ba8a4c6df775741732bd/llama/model.py#L76
With kv_heads:
https://github.com/facebookresearch/codellama/blob/427d6ac90f0b7db206bc4c62f4c5d38f92ca4d10/llama/model.py#L90
The text was updated successfully, but these errors were encountered: