search

horseee / LLM-Pruner

[NeurIPS 2023] LLM-Pruner: On the Structural Pruning of Large Language Models. Support Llama-3/3.1, Llama-2, LLaMA, BLOOM, Vicuna, Baichuan, TinyLlama, etc.
https://arxiv.org/abs/2305.11627
Apache License 2.0
815 stars 89 forks source link

Adaptation of GQA #64

Open junzhang-zj opened 1 month ago

junzhang-zj commented 1 month ago

Thank you for your solid work. I would like to ask if the current version is suitable for GQA architecture models, such as LLaMA-2-70B and LLaMA-3.

horseee commented 1 month ago

Hi,

Thanks for your interest in our paper. Currently not, we tested it and it has some bugs in the code.

We are currently working on the support of LLM-Pruner for Llama3 and Llama-3.1, and it would take some time before we can release it (Most likely in this week).

junzhang-zj commented 1 month ago

Thanks!

horseee commented 1 month ago

Hi,

The code supports GQA now. Here is the command and some of the results

junzhang-zj commented 1 month ago

Thanks, I'm using it now! But when I try param_mix on 70b, it gets aborted, while when I run param_first on the CPU, it works fine, albeit slower. I have another question, is there a specific reason for pruning based on k proj? Why not try to fix k and v, and prune q and o at a fixed ratio every 8 heads?

RamitPahwa commented 1 month ago

What are the memory requirements to run the code, I face OOM on 40GB A100, if I set my device to cuda instead of cpu , which is extremely slow, Is there a workaround to make it faster.

junzhang-zj commented 1 month ago

@RamitPahwa For the current version, it seems that for large models, normal pruning can only be done through the CPU.

horseee commented 1 month ago

What are the memory requirements to run the code, I face OOM on 40GB A100, if I set my device to cuda instead of cpu , which is extremely slow, Is there a workaround to make it faster.

Hi @RamitPahwa,

I met the same issue when I tried to prune llama-3. The pruning needs a GPU with 80GB of memory, and I'm not sure the reason for its extremely slow speed on CPUs. Perhaps some bugs in my code......