[shardformer] Support customized policy for llamav2 based model with HybridParallelPlugin

[eric8607242]

📌 Checklist before creating the PR

• I have created an issue for this PR for traceability
• The title follows the standard format: [doc/gemini/tensor/...]: A concise description
• I have added relevant tags if possible for us to better distinguish different PRs

🚨 Issue number

#4613

📝 What does this PR do?

In this PR, I have two modifications for ShardFormer.

1. I created a new argument policy for HybridParallelPlugin, which enables user can apply HybridParallel to their own model with customized Policy.
2. I add a new attribute replacement self_attn.num_key_value_heads for LlamaPolicy. The attribute is new for LLaMAv2, without this attribute replacement, I can not apply tensor parallelism on LLaMAv2 successfully.

💥 Checklist before requesting a review

• I have linked my PR to an issue (instruction)
• My issue clearly describes the problem/feature/proposal, with diagrams/charts/table/code if possible
• I have performed a self-review of my code
• I have added thorough tests.
• I have added docstrings for all the functions/methods I implemented

⭐️ Do you enjoy contributing to Colossal-AI?

• 🌝 Yes, I do.
• 🌚 No, I don't.

Tell us more if you don't enjoy contributing to Colossal-AI.

[flybird11111]

Thank you for your contribution. Could you please make some changes here? It's best not to expose 'Policy' in the plugin; it's just a component of Shardformer.

[eric8607242]

Hi,

Thanks for your review!

I have two candidate approaches currently.

1. Remove the typing hint
2. Modify the typing hint from Policy to Object

Which one do you think better?

[eric8607242]

Hi,

I modify the code with the second approach (modify from Policy to object)