github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2026-01-29 21:49:54 +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

11
tests/kit/model_zoo/transformers/chatglm.py
View File

@@ -3,7 +3,7 @@ import transformers
from ..registry import ModelAttribute, model_zoo
from .chatglm2_6b.configuration_chatglm import ChatGLMConfig
from .chatglm2_6b.modeling_chatglm import ChatGLMModel
from .chatglm2_6b.modeling_chatglm import ChatGLMForConditionalGeneration, ChatGLMModel
# ================================
# Register single-sentence ChatGLM
@@ -21,7 +21,7 @@ output_transform_fn = lambda x: x
# define loss function
loss_fn_for_chatglm_model = lambda x: x.last_hidden_state.mean()
loss_fn = lambda x: x.loss
loss_fn = lambda x: x.logits.mean()
config = ChatGLMConfig(num_layers=1,
padded_vocab_size=65024,
hidden_size=64,
@@ -36,3 +36,10 @@ model_zoo.register(name='transformers_chatglm',
output_transform_fn=output_transform_fn,
loss_fn=loss_fn_for_chatglm_model,
model_attribute=ModelAttribute(has_control_flow=True))
model_zoo.register(name="transformers_chatglm_for_conditional_generation",
model_fn=lambda: ChatGLMForConditionalGeneration(config, empty_init=False),
data_gen_fn=data_gen,
output_transform_fn=output_transform_fn,
loss_fn=loss_fn,
model_attribute=ModelAttribute(has_control_flow=True))