[shardformer] rename chatglm to chatglm2

colossalai/shardformer/policies/auto_policy.py

@@ -125,9 +125,9 @@ class PolicyLocation:
 
     # ChatGLM
     "colossalai.shardformer.modeling.chatglm2_6b.modeling_chatglm.ChatGLMModel":
-        PolicyLocation(file_name="chatglm", class_name="ChatGLMModelPolicy"),
+        PolicyLocation(file_name="chatglm2", class_name="ChatGLMModelPolicy"),
     "colossalai.shardformer.modeling.chatglm2_6b.modeling_chatglm.ChatGLMForConditionalGeneration":
-        PolicyLocation(file_name="chatglm", class_name="ChatGLMForConditionalGenerationPolicy"),
+        PolicyLocation(file_name="chatglm2", class_name="ChatGLMForConditionalGenerationPolicy"),
 }
 
 

colossalai/shardformer/policies/chatglm.py → colossalai/shardformer/policies/chatglm2.py

@@ -7,15 +7,15 @@
 
 import colossalai.shardformer.layer as col_nn
 from colossalai.pipeline.stage_manager import PipelineStageManager
-from colossalai.shardformer.modeling.chatglm import ChatGLMPipelineForwards
+from colossalai.shardformer.modeling.chatglm2 import ChatGLMPipelineForwards
 from colossalai.shardformer.modeling.chatglm2_6b.configuration_chatglm import ChatGLMConfig
 from colossalai.shardformer.modeling.chatglm2_6b.modeling_chatglm import (
     ChatGLMForConditionalGeneration,
     ChatGLMModel,
     GLMBlock,
 )
 
-from ..modeling.chatglm import get_flash_core_attention_forward, get_jit_fused_glm_block_forward
+from ..modeling.chatglm2 import get_flash_core_attention_forward, get_jit_fused_glm_block_forward
 from ..modeling.jit import get_jit_fused_dropout_add_func
 from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription
 

tests/kit/model_zoo/transformers/__init__.py

@@ -2,7 +2,7 @@
 from .bert import *
 from .blip2 import *
 from .bloom import *
-from .chatglm import *
+from .chatglm2 import *
 from .gpt import *
 from .llama import *
 from .opt import *