github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

to: add qwen2 auto policy

pull/5842/head
Wenhao Chen 2024-03-01 16:55:04 +08:00 committed by アマデウス
parent 6fa181ebef
commit 856b39f69d
1 changed files with 10 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
colossalai/shardformer/policies/auto_policy.py
View File

@ -182,6 +182,16 @@ _POLICY_LIST = {
"transformers.models.mistral.modeling_mistral.MistralForSequenceClassification": PolicyLocation(
file_name="mistral", class_name="MistralForSequenceClassificationPolicy"
),
# Qwen2
"transformers.models.qwen2.modeling_qwen2.Qwen2Model": PolicyLocation(
file_name="qwen2", class_name="Qwen2ModelPolicy"
),
"transformers.models.qwen2.modeling_qwen2.Qwen2ForCausalLM": PolicyLocation(
file_name="qwen2", class_name="Qwen2ForCausalLMPolicy"
),
"transformers.models.qwen2.modeling_qwen2.Qwen2ForSequenceClassification": PolicyLocation(
file_name="qwen2", class_name="Qwen2ForSequenceClassificationPolicy"
),
}