Fixed beam search error when using multiple GPUs

config.json

@@ -12,7 +12,7 @@
     "hidden_size": 5120,
     "initializer_range": 0.02,
     "intermediate_size": 16640,
-    "max_position_embeddings": 4096,
+    "max_position_embeddings": 8192,
     "model_type": "plamo",
     "n_shared_head": 8,
     "num_attention_heads": 40,

pytorch_model-00001-of-00003.bin

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

pytorch_model-00002-of-00003.bin

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

pytorch_model-00003-of-00003.bin

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

pytorch_model.bin.index.json

@@ -0,0 +1,370 @@
+{
+  "metadata": {
+    "total_size": 26199101440
+  },
+  "weight_map": {
+    "lm_head.weight": "pytorch_model-00003-of-00003.bin",
+    "model.embed_tokens.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.0.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.0.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.0.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.0.norm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.0.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.0.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.0.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.0.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.0.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.1.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.1.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.1.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.1.norm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.1.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.1.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.1.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.1.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.1.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.10.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.10.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.10.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.10.norm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.10.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.10.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.10.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.10.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.10.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.11.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.11.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.11.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.11.norm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.11.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.11.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.11.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.11.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.11.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.12.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.12.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.12.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.12.norm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.12.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.12.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.12.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.12.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.12.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.13.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.13.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.13.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.13.norm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.13.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.13.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.13.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.13.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.13.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.14.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.14.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.14.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.14.norm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.14.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.14.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.14.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.14.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.14.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.15.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.15.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.15.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.15.norm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.15.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.15.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.15.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.15.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.15.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.16.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.16.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.16.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.16.norm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.16.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.16.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.16.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.16.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.16.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.17.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.17.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.17.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.17.norm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.17.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.17.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.17.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.17.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.17.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.18.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.18.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.18.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.18.norm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.18.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.18.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.18.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.18.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.18.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.19.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.19.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.19.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.19.norm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.19.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.19.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.19.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.19.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.19.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.2.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.2.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.2.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.2.norm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.2.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.2.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.2.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.2.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.2.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.20.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.20.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.20.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.20.norm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.20.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.20.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.20.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.20.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.20.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.21.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.21.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.21.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.21.norm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.21.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.21.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.21.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.21.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.21.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.22.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.22.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.22.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.22.norm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.22.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.22.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.22.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.22.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.22.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.23.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.23.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.23.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.23.norm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.23.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.23.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.23.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.23.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.23.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.24.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.24.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.24.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.24.norm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.24.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.24.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.24.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.24.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.24.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.25.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.25.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.25.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.25.norm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.25.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.25.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.25.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.25.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.25.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.26.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.26.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.26.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.26.norm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.26.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.26.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.26.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.26.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.26.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.27.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.27.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.27.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.27.norm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.27.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.27.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.27.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.27.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.27.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.28.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.28.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.28.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.28.norm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.28.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.28.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.28.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.28.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.28.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.29.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.29.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.29.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.29.norm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.29.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.29.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.29.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.29.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.29.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.3.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.3.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.3.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.3.norm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.3.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.3.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.3.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.3.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.3.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.30.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.30.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.30.mlp.up_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.30.norm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.30.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.30.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.30.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.30.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.30.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.layers.31.mlp.down_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.31.mlp.gate_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.31.mlp.up_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.31.norm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.31.self_attn.k_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.31.self_attn.o_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.31.self_attn.q_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.31.self_attn.rotary_emb.inv_freq": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.31.self_attn.v_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.32.mlp.down_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.32.mlp.gate_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.32.mlp.up_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.32.norm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.32.self_attn.k_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.32.self_attn.o_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.32.self_attn.q_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.32.self_attn.rotary_emb.inv_freq": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.32.self_attn.v_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.33.mlp.down_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.33.mlp.gate_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.33.mlp.up_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.33.norm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.33.self_attn.k_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.33.self_attn.o_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.33.self_attn.q_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.33.self_attn.rotary_emb.inv_freq": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.33.self_attn.v_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.34.mlp.down_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.34.mlp.gate_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.34.mlp.up_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.34.norm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.34.self_attn.k_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.34.self_attn.o_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.34.self_attn.q_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.34.self_attn.rotary_emb.inv_freq": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.34.self_attn.v_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.35.mlp.down_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.35.mlp.gate_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.35.mlp.up_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.35.norm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.35.self_attn.k_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.35.self_attn.o_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.35.self_attn.q_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.35.self_attn.rotary_emb.inv_freq": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.35.self_attn.v_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.36.mlp.down_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.36.mlp.gate_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.36.mlp.up_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.36.norm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.36.self_attn.k_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.36.self_attn.o_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.36.self_attn.q_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.36.self_attn.rotary_emb.inv_freq": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.36.self_attn.v_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.37.mlp.down_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.37.mlp.gate_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.37.mlp.up_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.37.norm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.37.self_attn.k_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.37.self_attn.o_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.37.self_attn.q_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.37.self_attn.rotary_emb.inv_freq": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.37.self_attn.v_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.38.mlp.down_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.38.mlp.gate_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.38.mlp.up_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.38.norm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.38.self_attn.k_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.38.self_attn.o_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.38.self_attn.q_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.38.self_attn.rotary_emb.inv_freq": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.38.self_attn.v_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.39.mlp.down_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.39.mlp.gate_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.39.mlp.up_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.39.norm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.39.self_attn.k_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.39.self_attn.o_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.39.self_attn.q_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.39.self_attn.rotary_emb.inv_freq": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.39.self_attn.v_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.layers.4.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.4.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.4.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.4.norm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.4.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.4.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.4.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.4.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.4.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.5.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.5.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.5.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.5.norm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.5.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.5.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.5.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.5.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.5.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.6.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.6.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.6.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.6.norm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.6.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.6.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.6.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.6.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.6.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.7.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.7.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.7.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.7.norm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.7.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.7.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.7.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.7.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.7.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.8.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.8.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.8.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.8.norm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.8.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.8.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.8.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.8.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.8.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.9.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.9.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.9.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.9.norm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.9.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.9.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.9.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.9.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.layers.9.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.norm.weight": "pytorch_model-00003-of-00003.bin"
+  }
+}

tokenization_plamo.py

@@ -5,6 +5,7 @@ from shutil import copyfile
 from typing import Any, Dict, List, Optional, Tuple
 
 import sentencepiece as spm
+import transformers
 from transformers.tokenization_utils import PreTrainedTokenizer
 from transformers.utils import logging
 
@@ -34,12 +35,6 @@ class PlamoTokenizer(PreTrainedTokenizer):  # type: ignore
             kwargs["add_bos_token"] = False
         if "add_eos_token" not in kwargs:
             kwargs["add_eos_token"] = False
-        self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
-        self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
-        self.sp_model.Load(vocab_file)
-        self.vocab_file = vocab_file
-        self.add_bos_token = kwargs["add_bos_token"]
-        self.add_eos_token = kwargs["add_eos_token"]
 
         super().__init__(
             vocab_file=vocab_file,
@@ -55,6 +50,15 @@ class PlamoTokenizer(PreTrainedTokenizer):  # type: ignore
             **kwargs,
         )
 
+        self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
+        self.vocab_file = vocab_file
+        self.add_bos_token = kwargs["add_bos_token"]
+        self.add_eos_token = kwargs["add_eos_token"]
+        self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
+        self.sp_model.Load(vocab_file)
+
+        self.add_tokens(self.all_special_tokens_extended, special_tokens=True)
+
     # the functions below are copied from hf transformers LlamaTokenizer's implementation to fix the behaviour of the tokenizer
     # https://github.com/huggingface/transformers/blob/v4.30.2/src/transformers/models/llama/tokenization_llama.py
 
@@ -151,3 +155,7 @@ class PlamoTokenizer(PreTrainedTokenizer):  # type: ignore
                 fi.write(content_spiece_model)
 
         return (out_vocab_file,)
+
+
+class PlamoConfig(transformers.LlamaConfig):  # type: ignore
+    model_type = "plamo"