Use embed_tokens as output.weight

Author: mitmul

convert_hf_to_gguf.py

@@ -3597,7 +3597,18 @@ def set_gguf_parameters(self):
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
         del bid  # unused
 
-        if name.endswith(".dt_bias"):
+        if name.endswith(".embed_tokens.weight"):
+            # If there is no lm_head, we need to map the token embedding to the output layer
+            assert self.tensor_names is not None
+            if all(['lm_head' not in name for name in self.tensor_names]):
+                name_base = name.replace(".embed_tokens.weight", "")
+                output_name = "lm_head"
+
+                embed_tokens_mapped = self.map_tensor_name(name)
+                output_mapped = self.map_tensor_name(output_name) + ".weight"
+
+                return [(embed_tokens_mapped, data_torch), (output_mapped, data_torch)]
+        elif name.endswith(".dt_bias"):
             name = name.rpartition(".dt_bias")[0] + ".dt_proj.bias"
         elif name.endswith(".dt_norm_weight"):
             name = name.rpartition(".dt_norm_weight")[0] + ".dt_norm.weight"
@@ -3620,6 +3631,8 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter
             data_torch += 1.0
         elif name.endswith(".post_mlp_norm.weight"):
             data_torch += 1.0 / (5**1.5)
+        elif name.endswith(".norm.weight"):
+            data_torch += 1.0
         elif name.endswith(".gate_up_proj.weight"):
             # Split the combined gate_up tensor
             split_size = data_torch.shape[0] // 2

gguf-py/gguf/tensor_mapping.py

@@ -62,11 +62,10 @@ class TensorNameMap:
         # Output
         MODEL_TENSOR.OUTPUT: (
             "embed_out",                 # gptneox
-            "lm_head",                   # gpt2 mpt falcon llama-hf baichuan qwen mamba dbrx jais nemotron exaone olmoe olmo2 phimoe
+            "lm_head",                   # gpt2 mpt falcon llama-hf baichuan qwen mamba dbrx jais nemotron exaone olmoe olmo2 phimoe plamo2
             "output",                    # llama-pth bloom internlm2
             "word_embeddings_for_head",  # persimmon
             "lm_head.linear",            # phi2
-            "lm_head",                   # plamo2
             "output_layer",              # chatglm
             "head",                      # rwkv
             "head.out",                  # wavtokenizer

src/llama-model.cpp

@@ -956,6 +956,7 @@ void llama_model::load_hparams(llama_model_loader & ml) {
                 }
 
                 switch (hparams.n_layer) {
+                    case 16: type = LLM_TYPE_1B; break;
                     case 32: type = LLM_TYPE_8B; break;
                     default: type = LLM_TYPE_UNKNOWN;
                }
@@ -8233,6 +8234,9 @@ struct llm_build_plamo2 : public llm_graph_context {
             // ggml_graph_add_node(gf, model.layers[il].attn_norm);
             // cb(model.layers[il].attn_norm, "attn_norm", il);
 
+            ggml_graph_add_node(gf, model.layers[il].attn_norm);
+            cb(model.layers[il].attn_norm, "attn_norm_weight", il);
+
             // pre_mixer_norm
             cb(inpL, "attn_pre_norm_input", il);
             cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);