Add GGUF for starcoder2 (#34094)

* add starcoder2 arch support for gguf

* fix q6 test
This commit is contained in:
Vladislav Bronzov
2024-10-14 10:22:49 +02:00
committed by GitHub
parent 4c439173df
commit cb5ca3265f
3 changed files with 69 additions and 0 deletions

View File

@@ -54,6 +54,9 @@ class GgufIntegrationTests(unittest.TestCase):
gpt2_model_id = "mradermacher/gpt2-GGUF"
gpt2_original_model_id = "openai-community/gpt2"
gpt2_xl_model_id = "RichardErkhov/openai-community_-_gpt2-xl-gguf"
starcoder2_model_id = "QuantFactory/starcoder2-3b-GGUF"
starcoder2_fp16_model_id = "brittlewis12/starcoder2-3b-GGUF"
starcoder2_original_model_id = "bigcode/starcoder2-3b"
# standard quants
q4_0_gguf_model_id = "tinyllama-1.1b-chat-v1.0.Q4_0.gguf"
@@ -93,6 +96,8 @@ class GgufIntegrationTests(unittest.TestCase):
fp16_gpt2_model_id = "gpt2.f16.gguf"
q8_gpt2_model_id = "gpt2.Q8_0.gguf"
q6_k_gpt2_xl_model_id = "gpt2-xl.Q6_K.gguf"
q6_k_starcoder2_model_id = "starcoder2-3b.Q6_K.gguf"
fp16_starcoder2_gguf_model_id = "starcoder2-3b.fp16.gguf"
example_text = "Hello"
@@ -650,6 +655,45 @@ class GgufIntegrationTests(unittest.TestCase):
self.assertTrue(original_params.shape == converted_state_dict[layer_name].shape)
torch.testing.assert_close(original_params, converted_state_dict[layer_name])
def test_starcoder2_weights_conversion_fp16(self):
original_model = AutoModelForCausalLM.from_pretrained(
self.starcoder2_original_model_id,
device_map="auto",
torch_dtype=torch.float16,
)
converted_model = AutoModelForCausalLM.from_pretrained(
self.starcoder2_fp16_model_id,
gguf_file=self.fp16_starcoder2_gguf_model_id,
device_map="auto",
torch_dtype=torch.float16,
)
converted_state_dict = converted_model.state_dict()
original_state_dict = original_model.state_dict()
for layer_name, original_params in original_state_dict.items():
if layer_name in converted_state_dict and layer_name != "lm_head.weight":
# quantized models do not contain "lm_head.weight" layer
self.assertTrue(original_params.shape == converted_state_dict[layer_name].shape)
torch.testing.assert_close(original_params, converted_state_dict[layer_name])
def test_starcoder2_q6_k(self):
example_function_text = "def print_hello_world():"
model = AutoModelForCausalLM.from_pretrained(
self.starcoder2_model_id,
gguf_file=self.q6_k_starcoder2_model_id,
device_map="auto",
torch_dtype=torch.float16,
)
tokenizer = AutoTokenizer.from_pretrained(self.starcoder2_model_id, gguf_file=self.q6_k_starcoder2_model_id)
text = tokenizer(example_function_text, return_tensors="pt").to(torch_device)
out = model.generate(**text, max_new_tokens=10)
EXPECTED_TEXT = 'def print_hello_world():\n print("Hello World")\n\ndef print'
self.assertEqual(tokenizer.decode(out[0], skip_special_tokens=True), EXPECTED_TEXT)
def test_tokenization_xnli(self):
import tqdm
from datasets import load_dataset