[pipeline] add chatglm (#4363)

* add pipeline policy and bert forward to be done

* add bertmodel pipeline forward and make tests

* add Bert_Policy and test for policy

* update formatting

* update formatting

* update the code

* fix bugs

* fix name confilt

* add bloom model and policy ,revise the base class of policy

* revise

* revision

* add bert_for_pretraining

* add bert_for_pretraining forward and policy

* fix typos

* cancel warning

* change the imediate output to default dict

* change the default output of get_shared_params

* add chatglm

* add

* chatglm

* chatglm

* finish chatglm

* deletes

* fix rmsnorm

* chatglm

* fix chatglm shard

* init

tests/kit/model_zoo/transformers/chatglm.py

@@ -1,9 +1,11 @@
 import torch
 import transformers
 
+from colossalai.shardformer.modeling.chatglm2_6b.configuration_chatglm import ChatGLMConfig
+from colossalai.shardformer.modeling.chatglm2_6b.modeling_chatglm import ChatGLMForConditionalGeneration, ChatGLMModel
+
 from ..registry import ModelAttribute, model_zoo
-from .chatglm2_6b.configuration_chatglm import ChatGLMConfig
-from .chatglm2_6b.modeling_chatglm import ChatGLMForConditionalGeneration, ChatGLMModel
+
 
 # ================================
 # Register single-sentence ChatGLM
@@ -20,15 +22,18 @@ def data_gen():
 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.logits.mean()
+loss_fn_for_chatglm_model = lambda x: x.last_hidden_state.sum()
+loss_fn = lambda x: x.logits.sum()
+
 config = ChatGLMConfig(num_layers=1,
                        padded_vocab_size=65024,
                        hidden_size=64,
                        num_attention_heads=8,
-                       rmsnorm=False,
+                       rmsnorm=True,
                        original_rope=True,
-                       use_cache=True)
+                       use_cache=True,
+                       torch_dtype=torch.float32)
+
 
 model_zoo.register(name='transformers_chatglm',
                    model_fn=lambda: ChatGLMModel(config, empty_init=False),