starcoder2 (#317)

* add starcoder2 * upd * upd readme * upd readme * upd limits
smallcloudai · Feb 28, 2024 · eb84f4a · eb84f4a
1 parent a45af1e
commit eb84f4a
Show file tree

Hide file tree

Showing 4 changed files with 65 additions and 3 deletions.
diff --git a/README.md b/README.md
@@ -111,6 +111,9 @@ Extensions > Refact.ai Assistant > Settings > Infurl
 | [starcoder/7b/base](https://huggingface.co/smallcloudai/starcoderbase-7b)                         | +          |      | +           |
 | [starcoder/15b/base](https://huggingface.co/TheBloke/starcoder-GPTQ)                              | +          |      |             |
 | [starcoder/15b/plus](https://huggingface.co/TheBloke/starcoderplus-GPTQ)                          | +          |      |             |
+| [starcoder2/3b/base](https://huggingface.co/bigcode/starcoder2-3b)                                | +          |      | +           |
+| [starcoder2/7b/base](https://huggingface.co/bigcode/starcoder2-7b)                                | +          |      | +           |
+| [starcoder2/15b/base](https://huggingface.co/bigcode/starcoder2-15b)                              | +          |      | +           |
 | [wizardcoder/15b](https://huggingface.co/TheBloke/WizardCoder-15B-1.0-GPTQ)                       | +          |      |             |
 | [codellama/7b](https://huggingface.co/TheBloke/CodeLlama-7B-fp16)                                 | +          |      | +           |
 | [starchat/15b/beta](https://huggingface.co/TheBloke/starchat-beta-GPTQ)                           |            | +    |             |

diff --git a/known_models_db/refact_known_models/huggingface.py b/known_models_db/refact_known_models/huggingface.py
@@ -263,4 +263,34 @@
     #     "T": 4096,
     #     "filter_caps": ["chat"],
     # },
+    "starcoder2/3b/base": {
+        "backend": "transformers",
+        "model_path": "bigcode/starcoder2-3b",
+        "diff_scratchpad_class": "refact_scratchpads:ScratchpadPSM",
+        "chat_scratchpad_class": None,
+        "model_class_kwargs": {},
+        "required_memory_mb": 8000,
+        "T": 4096,
+        "filter_caps": ["completion", "finetune"],
+    },
+    "starcoder2/7b/base": {
+        "backend": "transformers",
+        "model_path": "bigcode/starcoder2-7b",
+        "diff_scratchpad_class": "refact_scratchpads:ScratchpadPSM",
+        "chat_scratchpad_class": None,
+        "model_class_kwargs": {},
+        "required_memory_mb": 16000,
+        "T": 4096,
+        "filter_caps": ["completion", "finetune"],
+    },
+    "starcoder2/15b/base": {
+        "backend": "transformers",
+        "model_path": "bigcode/starcoder2-15b",
+        "diff_scratchpad_class": "refact_scratchpads:ScratchpadPSM",
+        "chat_scratchpad_class": None,
+        "model_class_kwargs": {},
+        "required_memory_mb": 20000,
+        "T": 4096,
+        "filter_caps": ["completion", "finetune"],
+    }
 }
diff --git a/self_hosting_machinery/finetune/configuration/supported_models.py b/self_hosting_machinery/finetune/configuration/supported_models.py
@@ -40,7 +40,24 @@
     ],
     "force_enable_checkpointing": False
 }
-
+_starcoder2_base = {
+    "lora_target_modules_mapping": {
+        "qkv": ["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj"],
+        "out": ["self_attn.o_proj"],
+        "backproj": ["self_attn.o_proj"],
+        "mlp": ["mlp.c_fc", "mlp.c_proj"],
+    },
+    "freeze_exceptions_mapping": {
+        "wte": ["embed_tokens"],
+        "lm_head": ["lm_head"],
+        "lora": ["lora"]
+    },
+    "tokenizer": _bigcode_tokenizer_mapping,
+    "train_ds_pipeline": _fim_train_ds_pipeline,
+    "test_ds_pipeline": _fim_test_ds_pipeline,
+    "train_model_modifiers": [],
+    "force_enable_checkpointing": False
+}
 _deepseek_base = {
     "lora_target_modules_mapping": {
         "qkv": ["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj"],
@@ -103,6 +120,18 @@
         "force_enable_checkpointing": True
     },
 
+    "starcoder2/3b/base": _starcoder2_base,
+
+    "starcoder2/7b/base": {
+        **_starcoder2_base,
+        "force_enable_checkpointing": True
+    },
+
+    "starcoder2/15b/base": {
+        **_starcoder2_base,
+        "force_enable_checkpointing": True
+    },
+
     "codellama/7b": {
         "lora_target_modules_mapping": {
             "qkv": ["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj"],

diff --git a/setup.py b/setup.py
@@ -52,8 +52,8 @@ class PyPackage:
               "webgui/static/dashboards/*", "webgui/static/assets/*"]),
     "self_hosting_machinery": PyPackage(
         requires=["python-multipart", "auto-gptq==0.6.0", "accelerate",
-                  "termcolor", "torch", "transformers==4.38.1", "bitsandbytes",
-                  "safetensors", "peft", "triton",
+                  "termcolor", "torch", "transformers @ git+https://github.com/bigcode-project/transformers@starcoder2",
+                  "bitsandbytes", "safetensors", "peft", "triton",
                   "torchinfo", "mpi4py", "deepspeed==0.12.6",
                   "sentence-transformers", "huggingface-hub>=0.19.3"],
         optional=["ninja", "flash_attn @ git+https://github.com/smallcloudai/flash-attention@feat/alibi"],