search

princeton-nlp / LLM-Shearing

[ICLR 2024] Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning
https://arxiv.org/abs/2310.06694
MIT License
533 stars 39 forks source link

The Project is not implemented for 70B llama? #62

Open zhangzhenyu13 opened 6 months ago

zhangzhenyu13 commented 6 months ago

No GQA implementation is found, so the model is not capable to scale to 70B for composerLLAMA. Maybe we need design GQA and introduce head_z for wq and head_z_kv for wk and wv?

xiamengzhou commented 5 months ago

Hi, the modeling file currently does not support GQA, but should require minimal changes to support it. What you described should work perfectly :)

ZhiYuanZeng commented 5 months ago

It seems that we need a hierarchical pruning scheme for gqa, group pruning and head pruning inside group? Since we need to keep the number of heads in each group the same.

zhangzhenyu13 commented 5 months ago

It seems that we need a hierarchical pruning scheme for gqa, group pruning and head pruning inside group? Since we need to keep the number of heads in each group the same.

In order to make the pruned model be able to run tp, it would be better to keep the group num unchanged. We only need to prune the query heads for each group, thus maybe a layer_num * group_num* group_heads_query z_group_query need to initialized.

xiamengzhou commented 5 months ago

Pruning queries might cause the number of queries to be different in different groups. So maybe a group-based pruning is more reasonable? @zhangzhenyu13

ZhiYuanZeng commented 5 months ago

Could we share the mask of query-heads among different groups?

Pruning queries might cause the number of queries to be different in different groups. So maybe a group-based pruning is more reasonable? @zhangzhenyu13

zhangzhenyu13 commented 5 months ago

Pruning queries might cause the number of queries to be different in different groups. So maybe a group-based pruning is more reasonable? @zhangzhenyu13

Yes, your settings are right. We need to share z across groups.

Longyichen commented 5 months ago

Hi @zhangzhenyu13 I have some confusion. The author's composer llama file does not implement any GQA functionality. Did you implement GQA forward yourself? Which llama warehouse implementation version is better to refer to?