SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Add GGUF for Mamba (#34200)

Browse Source 
* add mamba architecture for gguf

* add logic for weights conversion, some fixes and refactoring

* add lm_head layers, unit test refactoring

* more fixes for tests

* remove lm_head creation

* remove unused comments
This commit is contained in:
Vladislav Bronzov
2024-10-30 16:52:17 +01:00
committed by GitHub
parent eab6c491d4
commit 5251fe6271
4 changed files with 93 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

56
tests/quantization/ggml/test_ggml.py
View File

@@ -59,6 +59,8 @@ class GgufIntegrationTests(unittest.TestCase):
starcoder2_model_id = "QuantFactory/starcoder2-3b-GGUF"
starcoder2_fp16_model_id = "brittlewis12/starcoder2-3b-GGUF"
starcoder2_original_model_id = "bigcode/starcoder2-3b"
mamba_original_model_id = "state-spaces/mamba-2.8b-hf"
mamba_model_id = "jpodivin/mamba-2.8b-hf-GGUF"
# standard quants
q4_0_gguf_model_id = "tinyllama-1.1b-chat-v1.0.Q4_0.gguf"
@@ -102,6 +104,8 @@ class GgufIntegrationTests(unittest.TestCase):
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"
q6_k_mamba_model_id = "ggml-model-Q6_K.gguf"
fp16_mamba_model_id = "ggml-model-f16.gguf"
example_text = "Hello"
@@ -573,6 +577,8 @@ class GgufIntegrationTests(unittest.TestCase):
if layer_name in quantized_state_dict:
self.assertTrue(original_params.shape == quantized_state_dict[layer_name].shape)
torch.testing.assert_close(original_params, quantized_state_dict[layer_name])
else:
raise ValueError(f"Layer {layer_name} is not presented in GGUF model")
def test_gpt2_xl_Q6_K(self):
tokenizer = AutoTokenizer.from_pretrained(self.gpt2_xl_model_id, gguf_file=self.q6_k_gpt2_xl_model_id)
@@ -639,6 +645,8 @@ class GgufIntegrationTests(unittest.TestCase):
if layer_name in quantized_state_dict:
self.assertTrue(original_params.shape == quantized_state_dict[layer_name].shape)
torch.testing.assert_close(original_params, quantized_state_dict[layer_name])
else:
raise ValueError(f"Layer {layer_name} is not presented in GGUF model")
def test_stablelm_q4_k_m(self):
model = AutoModelForCausalLM.from_pretrained(
@@ -708,6 +716,8 @@ class GgufIntegrationTests(unittest.TestCase):
if layer_name in converted_state_dict:
self.assertTrue(original_params.shape == converted_state_dict[layer_name].shape)
torch.testing.assert_close(original_params, converted_state_dict[layer_name])
else:
raise ValueError(f"Layer {layer_name} is not presented in GGUF model")
def test_starcoder2_weights_conversion_fp16(self):
original_model = AutoModelForCausalLM.from_pretrained(
@@ -727,10 +737,11 @@ class GgufIntegrationTests(unittest.TestCase):
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
if layer_name in converted_state_dict:
self.assertTrue(original_params.shape == converted_state_dict[layer_name].shape)
torch.testing.assert_close(original_params, converted_state_dict[layer_name])
else:
raise ValueError(f"Layer {layer_name} is not presented in GGUF model")
def test_starcoder2_q6_k(self):
example_function_text = "def print_hello_world():"
@@ -748,6 +759,47 @@ class GgufIntegrationTests(unittest.TestCase):
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_mamba_weights_conversion_fp16(self):
original_model = AutoModelForCausalLM.from_pretrained(
self.mamba_original_model_id,
torch_dtype=torch.float16,
)
converted_model = AutoModelForCausalLM.from_pretrained(
self.mamba_model_id,
gguf_file=self.fp16_mamba_model_id,
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:
self.assertTrue(original_params.shape == converted_state_dict[layer_name].shape)
if "mixer.A_log" in layer_name:
# we should increase tolerance after exponential reversing
# and performing np.log(-weights) operation as numbers are slightly different
torch.testing.assert_close(original_params, converted_state_dict[layer_name], atol=1e-3, rtol=1e-3)
else:
torch.testing.assert_close(original_params, converted_state_dict[layer_name])
else:
raise ValueError(f"Layer {layer_name} is not presented in GGUF model")
def test_mamba_q6_k(self):
model = AutoModelForCausalLM.from_pretrained(
self.mamba_model_id,
gguf_file=self.q6_k_mamba_model_id,
torch_dtype=torch.float16,
)
tokenizer = AutoTokenizer.from_pretrained(self.mamba_model_id, gguf_file=self.q6_k_mamba_model_id)
text = tokenizer(self.example_text, return_tensors="pt")["input_ids"]
out = model.generate(text, max_new_tokens=10)
EXPECTED_TEXT = "Hello,I answerthe question.\n\nA"
self.assertEqual(tokenizer.decode(out[0], skip_special_tokens=True), EXPECTED_TEXT)
def test_tokenization_xnli(self):
import tqdm
from datasets import load_dataset