HF transformers integration

config.json

@@ -1,9 +1,11 @@
 {
   "_name_or_path": "THUDM/chatglm3-6b",
-  "model_type": "chatglm",
+  "apply_query_key_layer_scaling": true,
   "architectures": [
-    "ChatGLMModel"
+    "ChatGlmForCausalLM"
   ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
   "auto_map": {
     "AutoConfig": "configuration_chatglm.ChatGLMConfig",
     "AutoModel": "modeling_chatglm.ChatGLMForConditionalGeneration",
@@ -11,32 +13,29 @@
     "AutoModelForSeq2SeqLM": "modeling_chatglm.ChatGLMForConditionalGeneration",
     "AutoModelForSequenceClassification": "modeling_chatglm.ChatGLMForSequenceClassification"
   },
-  "add_bias_linear": false,
-  "add_qkv_bias": true,
-  "apply_query_key_layer_scaling": true,
-  "apply_residual_connection_post_layernorm": false,
-  "attention_dropout": 0.0,
-  "attention_softmax_in_fp32": true,
-  "bias_dropout_fusion": true,
-  "ffn_hidden_size": 13696,
-  "fp32_residual_connection": false,
-  "hidden_dropout": 0.0,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "hidden_act": "silu",
   "hidden_size": 4096,
+  "initializer_range": 0.02,
+  "intermediate_size": 13696,
   "kv_channels": 128,
-  "layernorm_epsilon": 1e-05,
+  "max_position_embeddings": 2048,
+  "mlp_bias": false,
+  "model_type": "chatglm",
   "multi_query_attention": true,
   "multi_query_group_num": 2,
   "num_attention_heads": 32,
-  "num_layers": 28,
+  "num_hidden_layers": 28,
+  "num_key_value_heads": 32,
   "original_rope": true,
-  "padded_vocab_size": 65024,
-  "post_layer_norm": true,
-  "rmsnorm": true,
-  "seq_length": 8192,
-  "use_cache": true,
-  "torch_dtype": "float16",
-  "transformers_version": "4.30.2",
+  "partial_rotary_factor": 0.5,
+  "rms_norm_eps": 1e-05,
+  "rope_scaling": null,
+  "rope_theta": 10000.0,
   "tie_word_embeddings": false,
-  "eos_token_id": 2,
-  "pad_token_id": 0
-}
+  "torch_dtype": "float16",
+  "transformers_version": "4.37.0.dev0",
+  "use_cache": true,
+  "vocab_size": 65024
+}

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "transformers_version": "4.37.0.dev0"
+}

model-00001-of-00003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3684a0fd31c8a00e061d1242bcf0faadb34a9d0c70fb64d6ab40c703337e1cbe
+size 4907609888

model-00002-of-00003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:39afb798e4edce6865d67d65a32aa6fd9b47f545937aed82f16837146bc6bc59
+size 4895070096

model-00003-of-00003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ffa10cbc5159b8962f44210eb4ff6d180bfe867ae6f0c88615b0e4f58a0b0158
+size 2684511912

model.safetensors.index.json

@@ -1,207 +1,206 @@
 {
-    "metadata": {
-        "total_size": 12487168064
-    },
-    "weight_map": {
-        "transformer.embedding.word_embeddings.weight": "model-00001-of-00007.safetensors",
-        "transformer.encoder.final_layernorm.weight": "model-00007-of-00007.safetensors",
-        "transformer.encoder.layers.0.input_layernorm.weight": "model-00001-of-00007.safetensors",
-        "transformer.encoder.layers.0.mlp.dense_4h_to_h.weight": "model-00001-of-00007.safetensors",
-        "transformer.encoder.layers.0.mlp.dense_h_to_4h.weight": "model-00001-of-00007.safetensors",
-        "transformer.encoder.layers.0.post_attention_layernorm.weight": "model-00001-of-00007.safetensors",
-        "transformer.encoder.layers.0.self_attention.dense.weight": "model-00001-of-00007.safetensors",
-        "transformer.encoder.layers.0.self_attention.query_key_value.bias": "model-00001-of-00007.safetensors",
-        "transformer.encoder.layers.0.self_attention.query_key_value.weight": "model-00001-of-00007.safetensors",
-        "transformer.encoder.layers.1.input_layernorm.weight": "model-00001-of-00007.safetensors",
-        "transformer.encoder.layers.1.mlp.dense_4h_to_h.weight": "model-00001-of-00007.safetensors",
-        "transformer.encoder.layers.1.mlp.dense_h_to_4h.weight": "model-00001-of-00007.safetensors",
-        "transformer.encoder.layers.1.post_attention_layernorm.weight": "model-00001-of-00007.safetensors",
-        "transformer.encoder.layers.1.self_attention.dense.weight": "model-00001-of-00007.safetensors",
-        "transformer.encoder.layers.1.self_attention.query_key_value.bias": "model-00001-of-00007.safetensors",
-        "transformer.encoder.layers.1.self_attention.query_key_value.weight": "model-00001-of-00007.safetensors",
-        "transformer.encoder.layers.10.input_layernorm.weight": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.10.mlp.dense_4h_to_h.weight": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.10.mlp.dense_h_to_4h.weight": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.10.post_attention_layernorm.weight": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.10.self_attention.dense.weight": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.10.self_attention.query_key_value.bias": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.10.self_attention.query_key_value.weight": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.11.input_layernorm.weight": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.11.mlp.dense_4h_to_h.weight": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.11.mlp.dense_h_to_4h.weight": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.11.post_attention_layernorm.weight": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.11.self_attention.dense.weight": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.11.self_attention.query_key_value.bias": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.11.self_attention.query_key_value.weight": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.12.input_layernorm.weight": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.12.mlp.dense_4h_to_h.weight": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.12.mlp.dense_h_to_4h.weight": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.12.post_attention_layernorm.weight": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.12.self_attention.dense.weight": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.12.self_attention.query_key_value.bias": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.12.self_attention.query_key_value.weight": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.13.input_layernorm.weight": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.13.mlp.dense_4h_to_h.weight": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.13.mlp.dense_h_to_4h.weight": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.13.post_attention_layernorm.weight": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.13.self_attention.dense.weight": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.13.self_attention.query_key_value.bias": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.13.self_attention.query_key_value.weight": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.14.input_layernorm.weight": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.14.mlp.dense_4h_to_h.weight": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.14.mlp.dense_h_to_4h.weight": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.14.post_attention_layernorm.weight": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.14.self_attention.dense.weight": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.14.self_attention.query_key_value.bias": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.14.self_attention.query_key_value.weight": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.15.input_layernorm.weight": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.15.mlp.dense_4h_to_h.weight": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.15.mlp.dense_h_to_4h.weight": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.15.post_attention_layernorm.weight": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.15.self_attention.dense.weight": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.15.self_attention.query_key_value.bias": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.15.self_attention.query_key_value.weight": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.16.input_layernorm.weight": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.16.mlp.dense_4h_to_h.weight": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.16.mlp.dense_h_to_4h.weight": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.16.post_attention_layernorm.weight": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.16.self_attention.dense.weight": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.16.self_attention.query_key_value.bias": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.16.self_attention.query_key_value.weight": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.17.input_layernorm.weight": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.17.mlp.dense_4h_to_h.weight": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.17.mlp.dense_h_to_4h.weight": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.17.post_attention_layernorm.weight": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.17.self_attention.dense.weight": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.17.self_attention.query_key_value.bias": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.17.self_attention.query_key_value.weight": "model-00004-of-00007.safetensors",
-        "transformer.encoder.layers.18.input_layernorm.weight": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.18.mlp.dense_4h_to_h.weight": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.18.mlp.dense_h_to_4h.weight": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.18.post_attention_layernorm.weight": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.18.self_attention.dense.weight": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.18.self_attention.query_key_value.bias": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.18.self_attention.query_key_value.weight": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.19.input_layernorm.weight": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.19.mlp.dense_4h_to_h.weight": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.19.mlp.dense_h_to_4h.weight": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.19.post_attention_layernorm.weight": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.19.self_attention.dense.weight": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.19.self_attention.query_key_value.bias": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.19.self_attention.query_key_value.weight": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.2.input_layernorm.weight": "model-00001-of-00007.safetensors",
-        "transformer.encoder.layers.2.mlp.dense_4h_to_h.weight": "model-00001-of-00007.safetensors",
-        "transformer.encoder.layers.2.mlp.dense_h_to_4h.weight": "model-00001-of-00007.safetensors",
-        "transformer.encoder.layers.2.post_attention_layernorm.weight": "model-00001-of-00007.safetensors",
-        "transformer.encoder.layers.2.self_attention.dense.weight": "model-00001-of-00007.safetensors",
-        "transformer.encoder.layers.2.self_attention.query_key_value.bias": "model-00001-of-00007.safetensors",
-        "transformer.encoder.layers.2.self_attention.query_key_value.weight": "model-00001-of-00007.safetensors",
-        "transformer.encoder.layers.20.input_layernorm.weight": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.20.mlp.dense_4h_to_h.weight": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.20.mlp.dense_h_to_4h.weight": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.20.post_attention_layernorm.weight": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.20.self_attention.dense.weight": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.20.self_attention.query_key_value.bias": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.20.self_attention.query_key_value.weight": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.21.input_layernorm.weight": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.21.mlp.dense_4h_to_h.weight": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.21.mlp.dense_h_to_4h.weight": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.21.post_attention_layernorm.weight": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.21.self_attention.dense.weight": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.21.self_attention.query_key_value.bias": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.21.self_attention.query_key_value.weight": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.22.input_layernorm.weight": "model-00005-of-00007.safetensors",
-        "transformer.encoder.layers.22.mlp.dense_4h_to_h.weight": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.22.mlp.dense_h_to_4h.weight": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.22.post_attention_layernorm.weight": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.22.self_attention.dense.weight": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.22.self_attention.query_key_value.bias": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.22.self_attention.query_key_value.weight": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.23.input_layernorm.weight": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.23.mlp.dense_4h_to_h.weight": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.23.mlp.dense_h_to_4h.weight": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.23.post_attention_layernorm.weight": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.23.self_attention.dense.weight": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.23.self_attention.query_key_value.bias": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.23.self_attention.query_key_value.weight": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.24.input_layernorm.weight": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.24.mlp.dense_4h_to_h.weight": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.24.mlp.dense_h_to_4h.weight": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.24.post_attention_layernorm.weight": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.24.self_attention.dense.weight": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.24.self_attention.query_key_value.bias": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.24.self_attention.query_key_value.weight": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.25.input_layernorm.weight": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.25.mlp.dense_4h_to_h.weight": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.25.mlp.dense_h_to_4h.weight": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.25.post_attention_layernorm.weight": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.25.self_attention.dense.weight": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.25.self_attention.query_key_value.bias": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.25.self_attention.query_key_value.weight": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.26.input_layernorm.weight": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.26.mlp.dense_4h_to_h.weight": "model-00007-of-00007.safetensors",
-        "transformer.encoder.layers.26.mlp.dense_h_to_4h.weight": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.26.post_attention_layernorm.weight": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.26.self_attention.dense.weight": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.26.self_attention.query_key_value.bias": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.26.self_attention.query_key_value.weight": "model-00006-of-00007.safetensors",
-        "transformer.encoder.layers.27.input_layernorm.weight": "model-00007-of-00007.safetensors",
-        "transformer.encoder.layers.27.mlp.dense_4h_to_h.weight": "model-00007-of-00007.safetensors",
-        "transformer.encoder.layers.27.mlp.dense_h_to_4h.weight": "model-00007-of-00007.safetensors",
-        "transformer.encoder.layers.27.post_attention_layernorm.weight": "model-00007-of-00007.safetensors",
-        "transformer.encoder.layers.27.self_attention.dense.weight": "model-00007-of-00007.safetensors",
-        "transformer.encoder.layers.27.self_attention.query_key_value.bias": "model-00007-of-00007.safetensors",
-        "transformer.encoder.layers.27.self_attention.query_key_value.weight": "model-00007-of-00007.safetensors",
-        "transformer.encoder.layers.3.input_layernorm.weight": "model-00001-of-00007.safetensors",
-        "transformer.encoder.layers.3.mlp.dense_4h_to_h.weight": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.3.mlp.dense_h_to_4h.weight": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.3.post_attention_layernorm.weight": "model-00001-of-00007.safetensors",
-        "transformer.encoder.layers.3.self_attention.dense.weight": "model-00001-of-00007.safetensors",
-        "transformer.encoder.layers.3.self_attention.query_key_value.bias": "model-00001-of-00007.safetensors",
-        "transformer.encoder.layers.3.self_attention.query_key_value.weight": "model-00001-of-00007.safetensors",
-        "transformer.encoder.layers.4.input_layernorm.weight": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.4.mlp.dense_4h_to_h.weight": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.4.mlp.dense_h_to_4h.weight": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.4.post_attention_layernorm.weight": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.4.self_attention.dense.weight": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.4.self_attention.query_key_value.bias": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.4.self_attention.query_key_value.weight": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.5.input_layernorm.weight": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.5.mlp.dense_4h_to_h.weight": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.5.mlp.dense_h_to_4h.weight": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.5.post_attention_layernorm.weight": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.5.self_attention.dense.weight": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.5.self_attention.query_key_value.bias": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.5.self_attention.query_key_value.weight": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.6.input_layernorm.weight": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.6.mlp.dense_4h_to_h.weight": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.6.mlp.dense_h_to_4h.weight": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.6.post_attention_layernorm.weight": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.6.self_attention.dense.weight": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.6.self_attention.query_key_value.bias": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.6.self_attention.query_key_value.weight": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.7.input_layernorm.weight": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.7.mlp.dense_4h_to_h.weight": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.7.mlp.dense_h_to_4h.weight": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.7.post_attention_layernorm.weight": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.7.self_attention.dense.weight": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.7.self_attention.query_key_value.bias": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.7.self_attention.query_key_value.weight": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.8.input_layernorm.weight": "model-00002-of-00007.safetensors",
-        "transformer.encoder.layers.8.mlp.dense_4h_to_h.weight": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.8.mlp.dense_h_to_4h.weight": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.8.post_attention_layernorm.weight": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.8.self_attention.dense.weight": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.8.self_attention.query_key_value.bias": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.8.self_attention.query_key_value.weight": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.9.input_layernorm.weight": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.9.mlp.dense_4h_to_h.weight": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.9.mlp.dense_h_to_4h.weight": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.9.post_attention_layernorm.weight": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.9.self_attention.dense.weight": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.9.self_attention.query_key_value.bias": "model-00003-of-00007.safetensors",
-        "transformer.encoder.layers.9.self_attention.query_key_value.weight": "model-00003-of-00007.safetensors",
-        "transformer.output_layer.weight": "model-00007-of-00007.safetensors",
-        "transformer.rotary_pos_emb.inv_freq": "model-00001-of-00007.safetensors"
-    }
-}
+  "metadata": {
+    "total_size": 12487168000
+  },
+  "weight_map": {
+    "lm_head.weight": "model-00003-of-00003.safetensors",
+    "model.embed_tokens.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.mlp.dense_4h_to_h.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.mlp.dense_h_to_4h.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attention.dense.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attention.query_key_value.bias": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attention.query_key_value.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.mlp.dense_4h_to_h.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.mlp.dense_h_to_4h.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attention.dense.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attention.query_key_value.bias": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attention.query_key_value.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.10.mlp.dense_4h_to_h.weight": "model-00002-of-00003.safetensors",
+    "model.layers.10.mlp.dense_h_to_4h.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.10.self_attention.dense.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.self_attention.query_key_value.bias": "model-00001-of-00003.safetensors",
+    "model.layers.10.self_attention.query_key_value.weight": "model-00001-of-00003.safetensors",
+    "model.layers.11.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.11.mlp.dense_4h_to_h.weight": "model-00002-of-00003.safetensors",
+    "model.layers.11.mlp.dense_h_to_4h.weight": "model-00002-of-00003.safetensors",
+    "model.layers.11.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.11.self_attention.dense.weight": "model-00002-of-00003.safetensors",
+    "model.layers.11.self_attention.query_key_value.bias": "model-00002-of-00003.safetensors",
+    "model.layers.11.self_attention.query_key_value.weight": "model-00002-of-00003.safetensors",
+    "model.layers.12.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.12.mlp.dense_4h_to_h.weight": "model-00002-of-00003.safetensors",
+    "model.layers.12.mlp.dense_h_to_4h.weight": "model-00002-of-00003.safetensors",
+    "model.layers.12.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.12.self_attention.dense.weight": "model-00002-of-00003.safetensors",
+    "model.layers.12.self_attention.query_key_value.bias": "model-00002-of-00003.safetensors",
+    "model.layers.12.self_attention.query_key_value.weight": "model-00002-of-00003.safetensors",
+    "model.layers.13.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.13.mlp.dense_4h_to_h.weight": "model-00002-of-00003.safetensors",
+    "model.layers.13.mlp.dense_h_to_4h.weight": "model-00002-of-00003.safetensors",
+    "model.layers.13.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.13.self_attention.dense.weight": "model-00002-of-00003.safetensors",
+    "model.layers.13.self_attention.query_key_value.bias": "model-00002-of-00003.safetensors",
+    "model.layers.13.self_attention.query_key_value.weight": "model-00002-of-00003.safetensors",
+    "model.layers.14.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.14.mlp.dense_4h_to_h.weight": "model-00002-of-00003.safetensors",
+    "model.layers.14.mlp.dense_h_to_4h.weight": "model-00002-of-00003.safetensors",
+    "model.layers.14.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.14.self_attention.dense.weight": "model-00002-of-00003.safetensors",
+    "model.layers.14.self_attention.query_key_value.bias": "model-00002-of-00003.safetensors",
+    "model.layers.14.self_attention.query_key_value.weight": "model-00002-of-00003.safetensors",
+    "model.layers.15.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.15.mlp.dense_4h_to_h.weight": "model-00002-of-00003.safetensors",
+    "model.layers.15.mlp.dense_h_to_4h.weight": "model-00002-of-00003.safetensors",
+    "model.layers.15.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.15.self_attention.dense.weight": "model-00002-of-00003.safetensors",
+    "model.layers.15.self_attention.query_key_value.bias": "model-00002-of-00003.safetensors",
+    "model.layers.15.self_attention.query_key_value.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.mlp.dense_4h_to_h.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.mlp.dense_h_to_4h.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attention.dense.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attention.query_key_value.bias": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attention.query_key_value.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.mlp.dense_4h_to_h.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.mlp.dense_h_to_4h.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attention.dense.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attention.query_key_value.bias": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attention.query_key_value.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.mlp.dense_4h_to_h.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.mlp.dense_h_to_4h.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attention.dense.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attention.query_key_value.bias": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attention.query_key_value.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.mlp.dense_4h_to_h.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.mlp.dense_h_to_4h.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attention.dense.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attention.query_key_value.bias": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attention.query_key_value.weight": "model-00002-of-00003.safetensors",
+    "model.layers.2.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.mlp.dense_4h_to_h.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.mlp.dense_h_to_4h.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attention.dense.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attention.query_key_value.bias": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attention.query_key_value.weight": "model-00001-of-00003.safetensors",
+    "model.layers.20.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.mlp.dense_4h_to_h.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.mlp.dense_h_to_4h.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attention.dense.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attention.query_key_value.bias": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attention.query_key_value.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.mlp.dense_4h_to_h.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.mlp.dense_h_to_4h.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attention.dense.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attention.query_key_value.bias": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attention.query_key_value.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.22.mlp.dense_4h_to_h.weight": "model-00003-of-00003.safetensors",
+    "model.layers.22.mlp.dense_h_to_4h.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.22.self_attention.dense.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.self_attention.query_key_value.bias": "model-00002-of-00003.safetensors",
+    "model.layers.22.self_attention.query_key_value.weight": "model-00002-of-00003.safetensors",
+    "model.layers.23.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.23.mlp.dense_4h_to_h.weight": "model-00003-of-00003.safetensors",
+    "model.layers.23.mlp.dense_h_to_4h.weight": "model-00003-of-00003.safetensors",
+    "model.layers.23.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.23.self_attention.dense.weight": "model-00003-of-00003.safetensors",
+    "model.layers.23.self_attention.query_key_value.bias": "model-00003-of-00003.safetensors",
+    "model.layers.23.self_attention.query_key_value.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.mlp.dense_4h_to_h.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.mlp.dense_h_to_4h.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.self_attention.dense.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.self_attention.query_key_value.bias": "model-00003-of-00003.safetensors",
+    "model.layers.24.self_attention.query_key_value.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.mlp.dense_4h_to_h.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.mlp.dense_h_to_4h.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.self_attention.dense.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.self_attention.query_key_value.bias": "model-00003-of-00003.safetensors",
+    "model.layers.25.self_attention.query_key_value.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.mlp.dense_4h_to_h.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.mlp.dense_h_to_4h.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.self_attention.dense.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.self_attention.query_key_value.bias": "model-00003-of-00003.safetensors",
+    "model.layers.26.self_attention.query_key_value.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.mlp.dense_4h_to_h.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.mlp.dense_h_to_4h.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.self_attention.dense.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.self_attention.query_key_value.bias": "model-00003-of-00003.safetensors",
+    "model.layers.27.self_attention.query_key_value.weight": "model-00003-of-00003.safetensors",
+    "model.layers.3.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.mlp.dense_4h_to_h.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.mlp.dense_h_to_4h.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attention.dense.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attention.query_key_value.bias": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attention.query_key_value.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.mlp.dense_4h_to_h.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.mlp.dense_h_to_4h.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attention.dense.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attention.query_key_value.bias": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attention.query_key_value.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.mlp.dense_4h_to_h.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.mlp.dense_h_to_4h.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attention.dense.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attention.query_key_value.bias": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attention.query_key_value.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.mlp.dense_4h_to_h.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.mlp.dense_h_to_4h.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attention.dense.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attention.query_key_value.bias": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attention.query_key_value.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.mlp.dense_4h_to_h.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.mlp.dense_h_to_4h.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attention.dense.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attention.query_key_value.bias": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attention.query_key_value.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.mlp.dense_4h_to_h.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.mlp.dense_h_to_4h.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attention.dense.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attention.query_key_value.bias": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attention.query_key_value.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.mlp.dense_4h_to_h.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.mlp.dense_h_to_4h.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attention.dense.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attention.query_key_value.bias": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attention.query_key_value.weight": "model-00001-of-00003.safetensors",
+    "model.norm.weight": "model-00003-of-00003.safetensors"
+  }
+}

modeling_chatglm.py

@@ -223,8 +223,7 @@ class CoreAttention(torch.nn.Module):
         if pytorch_major_version >= 2:
             query_layer, key_layer, value_layer = [k.permute(1, 2, 0, 3) for k in [query_layer, key_layer, value_layer]]
             if attention_mask is None and query_layer.shape[2] == key_layer.shape[2]:
-                context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,
-                                                                                 is_causal=True)
+                context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,is_causal=True)
             else:
                 if attention_mask is not None:
                     attention_mask = ~attention_mask
@@ -312,7 +311,6 @@ class CoreAttention(torch.nn.Module):
 
 class SelfAttention(torch.nn.Module):
     """Parallel self-attention layer abstract class.
-
     Self-attention layer takes input with size [s, b, h]
     and returns output of the same size.
     """
@@ -448,7 +446,6 @@ class SelfAttention(torch.nn.Module):
 
         return output, kv_cache
 
-
 def _config_to_kwargs(args):
     common_kwargs = {
         "dtype": args.torch_dtype,
@@ -504,7 +501,6 @@ class MLP(torch.nn.Module):
 
 class GLMBlock(torch.nn.Module):
     """A single transformer layer.
-
     Transformer layer takes input with size [s, b, h] and returns an
     output of the same size.
     """
@@ -597,7 +593,7 @@ class GLMTransformer(torch.nn.Module):
         if self.post_layer_norm:
             LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm
             # Final layer norm before output.
-            self.final_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,
+            self.norm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,
                                                  dtype=config.torch_dtype)
 
         self.gradient_checkpointing = False
@@ -653,7 +649,7 @@ class GLMTransformer(torch.nn.Module):
 
         # Final layer norm.
         if self.post_layer_norm:
-            hidden_states = self.final_layernorm(hidden_states)
+            hidden_states = self.norm(hidden_states)
 
         return hidden_states, presents, all_hidden_states, all_self_attentions
 
@@ -740,7 +736,14 @@ class ChatGLMModel(ChatGLMPreTrainedModel):
         init_kwargs = {}
         if device is not None:
             init_kwargs["device"] = device
-        self.embedding = init_method(Embedding, config, **init_kwargs)
+
+        self.embed_tokens = nn.Embedding(
+            config.padded_vocab_size,
+            config.hidden_size,
+            dtype=config.torch_dtype,
+            device=device
+        )
+            
         self.num_layers = config.num_layers
         self.multi_query_group_num = config.multi_query_group_num
         self.kv_channels = config.kv_channels
@@ -753,9 +756,21 @@ class ChatGLMModel(ChatGLMPreTrainedModel):
 
         self.rotary_pos_emb = RotaryEmbedding(rotary_dim // 2, original_impl=config.original_rope, device=device,
                                               dtype=config.torch_dtype)
-        self.encoder = init_method(GLMTransformer, config, **init_kwargs)
-        self.output_layer = init_method(nn.Linear, config.hidden_size, config.padded_vocab_size, bias=False,
-                                        dtype=config.torch_dtype, **init_kwargs)
+
+        # Transformer layers.
+        def build_layer(layer_number):
+            return GLMBlock(config, layer_number, device=device)
+
+        self.layers = torch.nn.ModuleList([build_layer(i + 1) for i in range(self.num_layers)])
+        self.num_layers = config.num_layers
+        self.post_layer_norm = config.post_layer_norm
+        
+        if self.post_layer_norm:
+            LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm
+            # Final layer norm before output.
+            self.norm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,
+                                                 dtype=config.torch_dtype)
+
         self.pre_seq_len = config.pre_seq_len
         self.prefix_projection = config.prefix_projection
         if self.pre_seq_len is not None:
@@ -765,6 +780,8 @@ class ChatGLMModel(ChatGLMPreTrainedModel):
             self.prefix_encoder = PrefixEncoder(config)
             self.dropout = torch.nn.Dropout(0.1)
 
+        self.gradient_checkpointing = False
+
     def get_input_embeddings(self):
         return self.embedding.word_embeddings
 
@@ -804,7 +821,7 @@ class ChatGLMModel(ChatGLMPreTrainedModel):
         batch_size, seq_length = input_ids.shape
 
         if inputs_embeds is None:
-            inputs_embeds = self.embedding(input_ids)
+            inputs_embeds = self.embed_tokens(input_ids)
 
         if self.pre_seq_len is not None:
             if past_key_values is None:
@@ -827,10 +844,54 @@ class ChatGLMModel(ChatGLMPreTrainedModel):
         rotary_pos_emb = rotary_pos_emb.transpose(0, 1).contiguous()
 
         # Run encoder.
-        hidden_states, presents, all_hidden_states, all_self_attentions = self.encoder(
-            inputs_embeds, full_attention_mask, rotary_pos_emb=rotary_pos_emb,
-            kv_caches=past_key_values, use_cache=use_cache, output_hidden_states=output_hidden_states
-        )
+        if not past_key_values:
+            past_key_values = [None for _ in range(self.num_layers)]
+        presents = () if use_cache else None
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                )
+                use_cache = False
+
+        all_self_attentions = None
+        all_hidden_states = () if output_hidden_states else None
+
+        hidden_states = inputs_embeds
+        # To comply with former chat-glm format that expects (seqlen, bs, hd)
+        hidden_states = hidden_states.permute(1, 0, 2)
+        
+        for index, layer in enumerate(self.layers):
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                layer_ret = torch.utils.checkpoint.checkpoint(
+                    layer,
+                    hidden_states,
+                    full_attention_mask,
+                    rotary_pos_emb,
+                    past_key_values[index],
+                    use_cache
+                )
+            else:
+                layer_ret = layer(
+                    hidden_states,
+                    full_attention_mask,
+                    rotary_pos_emb,
+                    kv_cache=past_key_values[index],
+                    use_cache=use_cache
+                )
+            hidden_states, kv_cache = layer_ret
+            if use_cache:
+                presents = presents + (kv_cache,)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        # Final layer norm.
+        if self.post_layer_norm:
+            hidden_states = self.norm(hidden_states)
 
         if not return_dict:
             return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)
@@ -844,7 +905,7 @@ class ChatGLMModel(ChatGLMPreTrainedModel):
 
     def quantize(self, weight_bit_width: int):
         from .quantization import quantize
-        quantize(self.encoder, weight_bit_width)
+        quantize(self, weight_bit_width)
         return self
 
 
@@ -853,7 +914,8 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
         super().__init__(config)
 
         self.max_sequence_length = config.max_length
-        self.transformer = ChatGLMModel(config, empty_init=empty_init, device=device)
+        self.model = ChatGLMModel(config, empty_init=empty_init, device=device)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
         self.config = config
         self.quantized = False
 
@@ -934,7 +996,7 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
         use_cache = use_cache if use_cache is not None else self.config.use_cache
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
-        transformer_outputs = self.transformer(
+        transformer_outputs = self.model(
             input_ids=input_ids,
             position_ids=position_ids,
             attention_mask=attention_mask,
@@ -948,8 +1010,7 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
         hidden_states = transformer_outputs[0]
         if return_last_logit:
             hidden_states = hidden_states[-1:]
-        lm_logits = self.transformer.output_layer(hidden_states)
-        lm_logits = lm_logits.transpose(0, 1).contiguous()
+        lm_logits = self.lm_head(hidden_states)
 
         loss = None
         if labels is not None:
@@ -1062,8 +1123,8 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
         inputs = inputs.to(self.device)
         if past_key_values is not None:
             past_length = past_key_values[0][0].shape[0]
-            if self.transformer.pre_seq_len is not None:
-                past_length -= self.transformer.pre_seq_len
+            if self.model.pre_seq_len is not None:
+                past_length -= self.model.pre_seq_len
             inputs.position_ids += past_length
             attention_mask = inputs.attention_mask
             attention_mask = torch.cat((attention_mask.new_ones(1, past_length), attention_mask), dim=1)
@@ -1205,7 +1266,7 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
 
         self.config.quantization_bit = bits
 
-        self.transformer.encoder = quantize(self.transformer.encoder, bits, empty_init=empty_init, device=device,
+        self.model = quantize(self.model, bits, empty_init=empty_init, device=device,
                                             **kwargs)
         return self
 
@@ -1215,7 +1276,7 @@ class ChatGLMForSequenceClassification(ChatGLMPreTrainedModel):
         super().__init__(config)
 
         self.num_labels = config.num_labels
-        self.transformer = ChatGLMModel(config, empty_init=empty_init, device=device)
+        self.model = ChatGLMModel(config, empty_init=empty_init, device=device)
 
         self.classifier_head = nn.Linear(config.hidden_size, config.num_labels, bias=True, dtype=torch.half)
         if config.classifier_dropout is not None:
@@ -1242,7 +1303,7 @@ class ChatGLMForSequenceClassification(ChatGLMPreTrainedModel):
     ) -> Union[Tuple[torch.Tensor, ...], SequenceClassifierOutputWithPast]:
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
-        transformer_outputs = self.transformer(
+        transformer_outputs = self.model(
             input_ids=input_ids,
             position_ids=position_ids,
             attention_mask=attention_mask,
@@ -1293,4 +1354,4 @@ class ChatGLMForSequenceClassification(ChatGLMPreTrainedModel):
             past_key_values=transformer_outputs.past_key_values,
             hidden_states=transformer_outputs.hidden_states,
             attentions=transformer_outputs.attentions,
-        )
+        )