stablelm : StableLM support (ggerganov#3586)

* Add support for stablelm-3b-4e1t * Supports GPU offloading of (n-1) layers
Sanctum-AI · Nov 14, 2023 · 36eed0c · 36eed0c
1 parent b46d12f
commit 36eed0c
Show file tree

Hide file tree

Showing 6 changed files with 322 additions and 12 deletions.
diff --git a/README.md b/README.md
@@ -93,6 +93,7 @@ as the main playground for developing new features for the [ggml](https://github
 - [X] [Persimmon 8B](https://github.com/ggerganov/llama.cpp/pull/3410)
 - [X] [MPT](https://github.com/ggerganov/llama.cpp/pull/3417)
 - [X] [Bloom](https://github.com/ggerganov/llama.cpp/pull/3553)
+- [X] [StableLM-3b-4e1t](https://github.com/ggerganov/llama.cpp/pull/3586)
 
 
 **Bindings:**

diff --git a/convert-hf-to-gguf.py b/convert-hf-to-gguf.py
@@ -150,8 +150,6 @@ def load_hparams(dir_model):
 
     @staticmethod
     def from_model_architecture(model_architecture):
-        if model_architecture == "StableLMEpochForCausalLM":
-            return StableLMModel
         if model_architecture == "GPTNeoXForCausalLM":
             return GPTNeoXModel
         if model_architecture == "BloomForCausalLM":
@@ -168,6 +166,8 @@ def from_model_architecture(model_architecture):
             return RefactModel
         if model_architecture == "PersimmonForCausalLM":
             return PersimmonModel
+        if model_architecture in ("StableLMEpochForCausalLM", "LlavaStableLMEpochForCausalLM"):
+            return StableLMModel
         return Model
 
     def _is_model_safetensors(self) -> bool:
@@ -201,6 +201,8 @@ def _get_model_architecture(self) -> gguf.MODEL_ARCH:
             return gguf.MODEL_ARCH.REFACT
         if arch == "PersimmonForCausalLM":
             return gguf.MODEL_ARCH.PERSIMMON
+        if arch in ("StableLMEpochForCausalLM", "LlavaStableLMEpochForCausalLM"):
+            return gguf.MODEL_ARCH.STABLELM
 
         raise NotImplementedError(f'Architecture "{arch}" not supported!')
 
@@ -294,15 +296,6 @@ def _set_vocab_sentencepiece(self):
         special_vocab.add_to_gguf(self.gguf_writer)
 
 
-class StableLMModel(Model):
-    def set_gguf_parameters(self):
-        super().set_gguf_parameters()
-        self.gguf_writer.add_rope_dimension_count(
-            int(self.hparams["rope_pct"] * (self.hparams["hidden_size"] // self.hparams["num_attention_heads"])),
-        )
-        self.gguf_writer.add_layer_norm_eps(1e-5)
-
-
 class GPTNeoXModel(Model):
     def set_gguf_parameters(self):
         block_count = self.hparams["num_hidden_layers"]
@@ -824,6 +817,21 @@ def write_tensors(self):
             self.gguf_writer.add_tensor(new_name, data)
 
 
+class StableLMModel(Model):
+    def set_gguf_parameters(self):
+        hparams = self.hparams
+        block_count = hparams["num_hidden_layers"]
+
+        self.gguf_writer.add_name(dir_model.name)
+        self.gguf_writer.add_context_length(hparams["max_position_embeddings"])
+        self.gguf_writer.add_embedding_length(hparams["hidden_size"])
+        self.gguf_writer.add_block_count(block_count)
+        self.gguf_writer.add_feed_forward_length(hparams["intermediate_size"])
+        self.gguf_writer.add_rope_dimension_count(int(hparams["rope_pct"]*(hparams["hidden_size"] // hparams["num_attention_heads"])))
+        self.gguf_writer.add_head_count(hparams["num_attention_heads"])
+        self.gguf_writer.add_parallel_residual(hparams["use_parallel_residual"] if "use_parallel_residual" in hparams else True)
+        self.gguf_writer.add_layer_norm_eps(1e-5)
+
 ###### CONVERSION LOGIC ######
 
 def parse_args() -> argparse.Namespace:

diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py
@@ -90,6 +90,7 @@ class MODEL_ARCH(IntEnum):
     REFACT    = auto()
     BERT      = auto()
     BLOOM     = auto()
+    STABLELM  = auto()
 
 
 class MODEL_TENSOR(IntEnum):
@@ -129,6 +130,7 @@ class MODEL_TENSOR(IntEnum):
     MODEL_ARCH.REFACT:         "refact",
     MODEL_ARCH.BERT:           "bert",
     MODEL_ARCH.BLOOM:          "bloom",
+    MODEL_ARCH.STABLELM:       "stablelm",
 }
 
 TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
@@ -299,6 +301,21 @@ class MODEL_TENSOR(IntEnum):
         MODEL_TENSOR.FFN_DOWN,
         MODEL_TENSOR.FFN_UP,
     ],
+    MODEL_ARCH.STABLELM: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ROPE_FREQS,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
     MODEL_ARCH.GPT2: [
         # TODO
     ],