github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-25 11:44:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

[shardformer] support ChatGLMForConditionalGeneration & add fusedlayernorm for vit

Browse Source
This commit is contained in:
klhhhhh
2023-07-20 19:14:04 +08:00
committed by Hongxin Liu
parent 4da05052f4
commit 8120eca0c0
4 changed files with 37 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
colossalai/shardformer/policies/chatglm.py
View File

@@ -90,7 +90,31 @@ class ChatGLMModelPolicy(Policy):
policy=policy,
target_key=ChatGLMModel)
else:
self.append_or_create_submodule_replacement(description=[
SubModuleReplacementDescription(suffix="input_layernorm", target_module=col_nn.FusedRMSNorm),
SubModuleReplacementDescription(suffix="post_attention_layernorm",
target_module=col_nn.FusedRMSNorm)
],
policy=policy,
target_key=GLMBlock)
if self.model.config.post_layer_norm:
self.append_or_create_submodule_replacement(description=[
SubModuleReplacementDescription(suffix="encoder.final_layernorm",
target_module=col_nn.FusedRMSNorm)
],
policy=policy,
target_key=ChatGLMModel)
return policy
def postprocess(self):
return self.model
class ChatGLMForConditionalGenerationPolicy(ChatGLMModelPolicy):
def module_policy(self):
policy = super().module_policy()
return policy

2
colossalai/shardformer/policies/vit.py
View File

@@ -23,7 +23,7 @@ class ViTPolicy(Policy):
return self.model
def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]:
from transformers.models.vit.modeling_vit import ViTEmbeddings, ViTLayer
from transformers.models.vit.modeling_vit import ViTEmbeddings, ViTLayer, ViTModel
policy = {}