[BUG]: HybridParallelPlugin does not support customized wrapped LLaMAv2 model

[eric8607242]

🐛 Describe the bug

Hi,
in my use case, I have to wrap LLaMAv2 with an additional class.

There are two issues when I try to apply HybridParallelPlugin on my customized model:

1. HybridParallelPlugin can not assign the customized policy for ShardFormer, which makes HybridParallePlugin incompatible with the customized model.
2. Policy LlamaPolicy lacks the attribute self_attn.num_key_value_heads, which is a new property for LLaMAv2.

Environment

No response

[Issues-translate-bot]

Bot detected the issue body's language is not English, translate it automatically. 👯👭🏻🧑‍🤝‍🧑👫🧑🏿‍🤝‍🧑🏻👩🏾‍🤝‍👨🏿👬🏿

---

Title: [BUG]: Hybrid Parallel Plugin does not support custom wrapped LLaMA 2 model