Can the attention part of llama be searched?

[shijian971]

as follows
nncf_lonas_llama_7b_config.json

[Yuan0320]

Hi @shijian971 Yes, you just need to modify the model code (modeling_llama.py) manually. Please refer to the following modifications made to BERT: https://github.com/IntelLabs/Hardware-Aware-Automated-Machine-Learning/blob/main/LoNAS/patches/transformers-v4.31.0.patch#L43-L62. The main goal is to make num_attention_heads dynamic.

Btw, pruning attention heads can be a little sensitive for large models and complex tasks, which is why we have not pruned attention heads in some settings. We recommend starting with a smaller search space for attention part. Thanks! 😊