Views
No views yet
batch_size is set quite high as the model is small, you may need to adjust this to your GPU VRAM.1from gptqmodel import GPTQModel
2from transformers import AutoTokenizer
3
4# Use the local directory or JustJaro/Rombos-LLM-V2.6-Qwen-14b_gptq_g32_4bit after upload
5quantized_model_id = "/home/jaro/models/quantized/Rombos-LLM-V2.6-Qwen-14b_gptq_g32_4bit" # or "JustJaro/Rombos-LLM-V2.6-Qwen-14b_gptq_g32_4bit"
6tokenizer = AutoTokenizer.from_pretrained(quantized_model_id)
7model = GPTQModel.load(quantized_model_id, device="cuda:0") # or "cpu"
8
9input_text = "This is a test prompt"
10inputs = tokenizer(input_text, return_tensors="pt").to("cuda:0")
11outputs = model.generate(**inputs)
12print(tokenizer.decode(outputs[0], skip_special_tokens=True))1uv venv
2source venv/bin/activate
3uv sync1HF_TOKEN=<YOUR_HF_TOKEN>
2TOKENIZERS_PARALLELISM="true"
3PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True1#!/usr/bin/env python3
2"""
3This script loads a source Hugging Face model and a calibration dataset,
4quantizes the model using GPTQModel (with 4-bit precision and group size 128),
5saves the quantized model using the Transformers API with safetensors (safe serialization)
6under ~/models/quantized/, and then creates/updates a Hugging Face repository (with the
7_gptq_g128_4bit suffix) by uploading the model, tokenizer, and an auto-generated README.md.
8
9Usage example:
10 python quantize.py --source-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 \
11 --calibration-dataset wikitext/wikitext-2-raw-v1 \
12 --seq-len 1024 --nsamples 256 --hf-token <YOUR_HF_TOKEN>
13"""
14
15import os
16import shutil
17import subprocess
18from enum import Enum
19from pathlib import Path
20from typing import List
21
22import torch
23import typer
24from datasets import load_dataset
25from dotenv import load_dotenv, find_dotenv
26from gptqmodel import GPTQModel, QuantizeConfig
27from gptqmodel.utils import Perplexity
28# For later pushing to the model hub
29from huggingface_hub import HfApi
30from transformers import AutoTokenizer, PreTrainedTokenizerBase
31
32load_dotenv(find_dotenv())
33HF_TOKEN = os.getenv("HF_TOKEN")
34
35app = typer.Typer()
36
37class GroupSize(str, Enum):
38 accurate:int = 32
39 balanced:int = 64
40 fast:int = 128
41
42
43def get_text_from_example(example: dict) -> str:
44 """
45 Returns text from a dataset example.
46 If the example contains a "text" field, and it is nonempty, that text is used.
47 Otherwise, if it has a "messages" field (a list of dicts with a "content" key),
48 the function returns the concatenation of all non-empty message contents.
49 """
50 if "text" in example and example["text"]:
51 return example["text"]
52 elif "messages" in example:
53 contents = [msg.get("content", "").strip() for msg in example["messages"]]
54 return " ".join([s for s in contents if s])
55 else:
56 return ""
57
58
59def get_calibration_dataset(
60 tokenizer: PreTrainedTokenizerBase,
61 nsamples: int,
62 seqlen: int,
63 calibration_dataset: str
64 ) -> List[dict]:
65 """
66 Loads a calibration dataset from the Hugging Face Hub (or from a local file).
67 It accepts datasets with a single "text" field (like wikitext)
68 or with a "messages" field (as in the Neural Magic LLM Compression Calibration dataset).
69 Only examples whose extracted text length is at least 'seqlen' are kept.
70 Each chosen example is tokenized (with truncation up to 'seqlen') and returned as a dict.
71 """
72 ds = None
73 try:
74 # Attempt to load from HF Hub.
75 try:
76 if "/" in calibration_dataset:
77 parts = calibration_dataset.split("/", 1)
78 ds = load_dataset(parts[0], parts[1], split="train")
79 else:
80 ds = load_dataset(calibration_dataset, split="train")
81 except Exception as e:
82 print(f"Error loading dataset '{calibration_dataset}' via load_dataset: {e}")
83 ds = load_dataset(calibration_dataset, split="train")
84 print(f"Loaded calibration dataset from full remote path {calibration_dataset}.")
85
86
87 except Exception as e:
88 print(f"Error loading dataset '{calibration_dataset}' via load_dataset: {e}")
89 # Fallback: if the supplied calibration_dataset is a local path, try to load it as JSON-lines.
90 if os.path.exists(calibration_dataset):
91 try:
92 ds = load_dataset("json", data_files=calibration_dataset, split="train")
93 print(f"Loaded calibration dataset from local file {calibration_dataset}.")
94 except Exception as e2:
95 print(f"Error loading local json dataset from '{calibration_dataset}': {e2}")
96 return []
97 else:
98 return []
99
100 print(f"Dataset features: {ds.features}")
101
102 # Filter examples that have at least 80% 'seqlen' of extracted text (wikitext-2-raw-v1 dataset has short examples).
103 ds = ds.filter(lambda x: len(get_text_from_example(x)) <= int(seqlen*0.8))
104 sample_range = min(nsamples, len(ds))
105 calibration_data = []
106 for i in range(sample_range):
107 example = ds[i]
108 text = get_text_from_example(example)
109 tokenized = tokenizer(text, truncation=True, max_length=seqlen, return_tensors="pt")
110 tokenized = {k: v.squeeze(0) for k, v in tokenized.items()}
111 calibration_data.append(tokenized)
112 return calibration_data
113
114
115def calculate_avg_ppl(model, tokenizer):
116 """
117 Computes the average perplexity on the wikitext-2-raw-v1 train split using GPTQModel's Perplexity utility.
118 """
119 ppl = Perplexity(
120 model=model,
121 tokenizer=tokenizer,
122 dataset_path="wikitext",
123 dataset_name="wikitext-2-raw-v1",
124 split="train",
125 text_column="text",
126 )
127 ppl_values = ppl.calculate(n_ctx=512, n_batch=512)
128 avg = sum(ppl_values) / len(ppl_values)
129 return avg
130
131
132def get_pinned_package_versions():
133 """
134 Retrieves pinned package versions using 'uv pip freeze'.
135 Returns a dictionary mapping lowercased package names to their versions.
136 """
137 try:
138 result = subprocess.run(["uv", "pip", "freeze"], capture_output=True, text=True, check=True)
139 packages_output = result.stdout.strip()
140 versions = {}
141 for line in packages_output.splitlines():
142 if "==" in line:
143 package_name, package_version = line.split("==", 1)
144 versions[package_name.lower()] = package_version
145 return versions
146 except subprocess.CalledProcessError as e:
147 typer.echo(f"Error running 'uv pip freeze': {e}", err=True)
148 return {}
149 except FileNotFoundError:
150 typer.echo("uv command not found. Make sure uv is installed and in your PATH.", err=True)
151 return {}
152
153
154@app.command()
155def main(
156 seq_len: int = typer.Option(4096, help="Sequence length for tokenization and calibration."),
157 nsamples: int = typer.Option(512, help="Number of samples to use for calibration."),
158 source_model: str = typer.Option("rombodawg/Rombos-LLM-V2.6-Qwen-14b",
159 help="Source model HF repository identifier."),
160 calibration_dataset: str = typer.Option("wikitext/wikitext-2-raw-v1",
161 help="Calibration dataset identifier (in 'dataset/config' format) or local file path."),
162 hf_token: str = typer.Option(HF_TOKEN,
163 help="Hugging Face token for creating/updating your repo."),
164 upload_only: bool = typer.Option(False, help="Only upload the quantized model to the Hugging Face Hub."),
165 # Allow for 32, 64, 128 only using typer:
166 group_size: GroupSize = typer.Option(GroupSize.accurate, help="Group size for quantization accurate: 32, "
167 "balanced: 64, fast: 128. Default: accurate."),
168 mse: bool = typer.Option(True, help="Use mse instead of mae for the loss function."),
169 size_multi: int = typer.Option(3.5, help="Model size multiplier depends on the source model. Default: 1."),
170):
171 # Prepare destination directory and model names.
172 model_name = source_model.split("/")[-1]
173 if not size_multi == 1:
174 size_multiplier = size_multi
175 size_multiplier_len = size_multiplier / 2
176 else:
177 size_multiplier = 1
178 size_multiplier_len = 1
179 nsamples = int(nsamples * size_multiplier)
180 seq_len = int(seq_len * size_multiplier_len)
181 quantized_model_name = f"{model_name}_gptq_g{int(group_size.value)}_4bit"
182 quantized_model_dir = os.path.expanduser(os.path.join("~/models/quantized", quantized_model_name))
183 if not upload_only:
184 # Remove the directory if it already exists
185 if os.path.exists(quantized_model_dir):
186 shutil.rmtree(quantized_model_dir)
187 # Create directory for quantized model.
188 os.makedirs(quantized_model_dir, exist_ok=True)
189
190 typer.echo("Loading tokenizer from source model...")
191 tokenizer_obj = AutoTokenizer.from_pretrained(source_model, use_fast=True)
192
193 typer.echo("Loading calibration dataset...")
194 typer.echo(f"Calibration dataset: {calibration_dataset}")
195 calibration_data = get_calibration_dataset(tokenizer_obj, nsamples, seq_len, calibration_dataset)
196 if not calibration_data:
197 typer.echo("Calibration dataset is empty. Aborting.", err=True)
198 raise typer.Exit(code=1)
199 if mse:
200 # Fits mistral-small-24b particularly well, as well as the increased damp_percent
201 mse = 0.01
202 quantize_config = QuantizeConfig(bits=4, group_size=int(group_size.value), damp_percent=0.015, mse=mse)
203 else:
204 quantize_config = QuantizeConfig(bits=4, group_size=int(group_size.value), damp_percent=0.01)
205 device = "cuda:0" if torch.cuda.is_available() else "cpu"
206 typer.echo(f"Loading model in {device} mode...")
207 model = GPTQModel.load(source_model, quantize_config)
208
209 typer.echo("Quantizing model...")
210 group_size_factor = int(128 / int(group_size.value))
211 model.quantize(calibration_data, auto_gc=False,
212 batch_size=max(1, int(int((nsamples * 0.1) / group_size_factor) *
213 int(size_multiplier_len))))
214 # Retrieve Hugging Face user info for README generation.
215 package_versions = get_pinned_package_versions()
216 username = get_my_user(hf_token)
217
218 script_content = self_read_script()
219
220 typer.echo(f"Saving quantized model to {quantized_model_dir} using Transformers safe serialization...")
221 try:
222 model.save_pretrained(quantized_model_dir)
223 tokenizer_obj.save_pretrained(quantized_model_dir)
224 except Exception as ex:
225 typer.echo(f"Error during saving with safe_serialization: {ex}. Aborting.")
226 raise
227 typer.echo(f"Model uploaded to Hugging Face repo: {quantized_model_name}")
228 else:
229 tokenizer_obj = AutoTokenizer.from_pretrained(source_model, use_fast=True)
230 package_versions = get_pinned_package_versions()
231 username = get_my_user(hf_token)
232 script_content = self_read_script()
233
234
235 device = "cuda:0" if torch.cuda.is_available() else "cpu"
236 model = GPTQModel.load(quantized_model_dir, device=device)
237 avg_ppl = calculate_avg_ppl(model, tokenizer_obj)
238 typer.echo(f"Average perplexity (PPL) on wikitext v2 dataset: {avg_ppl}")
239 deps = Path("./pyproject.toml")
240 shutil.copy(deps, quantized_model_dir)
241 generate_readme(calibration_dataset, nsamples, quantized_model_dir,
242 quantized_model_name, script_content, seq_len, source_model, username, avg_ppl)
243 GPTQModel.push_to_hub(quantized_path=quantized_model_dir, private=False, repo_id=quantized_model_name,
244 token=HF_TOKEN)
245 typer.echo(f"Model uploaded to Hugging Face repo: {quantized_model_name}")
246 demo_input = tokenizer_obj("test is", return_tensors="pt").to(device)
247 generated_ids = model.generate(**demo_input)
248 output_text = tokenizer_obj.decode(generated_ids[0])
249 typer.echo(f"Inference demo output: {output_text}")
250 typer.echo(f"Average perplexity (PPL) on calibration dataset: {avg_ppl}")
251
252
253def self_read_script():
254 try:
255 script_path = os.path.abspath(__file__)
256 with open(script_path, "r") as f:
257 script_content = f.read()
258 except Exception as e:
259 script_content = "Error reading script content: " + str(e)
260 return script_content
261
262
263def get_my_user(hf_token):
264 api = HfApi(token=hf_token)
265 user_info = api.whoami()
266 try:
267 username = user_info.get("name") or user_info.get("username")
268 except Exception as e:
269 typer.echo(f"Error retrieving username from Hugging Face API: {e}. Using default username.")
270 username = api.whoami()
271 if not username:
272 typer.echo("Could not determine your Hugging Face username from the token, defaulting to hard coded username.",
273 err=True)
274 username = "JustJaro"
275 return username
276
277
278def generate_readme(calibration_dataset, nsamples, quantized_model_dir,
279 quantized_model_name, script_content, seq_len, source_model, username, avg_ppl):
280 readme_content = f"""---
281tags:
282- gptq
283- quantization
284- 4bit
285- confidentialmind
286- text-generation
287- apache2.0
288- mistral-small-24b
289---
290# 🔥 Quantized Model: {quantized_model_name} 🔥
291
292This is a 4-bit quantized version of [{source_model}](https://huggingface.co/{source_model}) model, quantized by [ConfidentialMind.com](https://www.confidentialmind.com) 🤖✨
293It leverages the open-source GPTQModel quantization to achieve 4-bit precision with a group size of 128 resulting in a
294smaller,
295faster model with minimal performance degradation.
296
297Ran on a single NVIDIA A100 GPU with 80GB of VRAM.
298
299*Note* `batch_size` is set quite high as the model is small, you may need to adjust this to your GPU VRAM.
300
301## Model Details
302- **Original Model:** [{source_model}](https://huggingface.co/{source_model})
303- **Quantized Model:** {quantized_model_name} (this repository)
304- **Quantization Method:** GPTQ (4-bit, group size 128)
305- **Quantization Library:** [GPTQModel](https://github.com/ModelCloud/GPTQModel/tree/main)
306- **Calibration Dataset:** {calibration_dataset} (using {nsamples} samples with seq len {seq_len})
307- **Quantized by:** [ConfidentialMind.com](https://www.confidentialmind.com)
308
309## Usage
310
311```python
312from gptqmodel import GPTQModel
313from transformers import AutoTokenizer
314
315# Use the local directory or {username}/{quantized_model_name} after upload
316quantized_model_id = "{quantized_model_dir}" # or "{username}/{quantized_model_name}"
317tokenizer = AutoTokenizer.from_pretrained(quantized_model_id)
318model = GPTQModel.load(quantized_model_id, device="cuda:0") # or "cpu"
319
320input_text = "This is a test prompt"
321inputs = tokenizer(input_text, return_tensors="pt").to("cuda:0")
322outputs = model.generate(**inputs)
323print(tokenizer.decode(outputs[0], skip_special_tokens=True))1uv venv
2source venv/bin/activate
3uv sync1HF_TOKEN=<YOUR_HF_TOKEN>
2TOKENIZERS_PARALLELISM="true"
3PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True{script_content}{source_model}gptqmodel.utils.eval integration and auto-generation of eval table.
## Quantization Performance
Average perplexity (PPL) on wikitext v2 dataset: 108.12590932665465
## Disclaimer
This model is for research purposes only. It may inherit limitations and biases from the original model and the quantization process. Please use responsibly and refer to the original model card for more details.
## Contact
For any questions or support, please visit [ConfidentialMind.com](https://www.confidentialmind.com) or contact us directly.
## License
This model inherits the license from the original model. Please refer to the original model card for more details.
Original model card: `rombodawg/Rombos-LLM-V2.6-Qwen-14b`
## Author
This model was quantized by [Jaro](https://www.linkedin.com/in/jaroai/)
## Acknowledgements
Quantization performed using the GPTQModel pipeline.
TODO: Add `gptqmodel.utils.eval` integration and auto-generation of eval table.
---
*Generated and quantized using GPTQModel.*