diff --git a/quantize_model.py b/quantize_model.py new file mode 100644 index 0000000..c1d1f3c --- /dev/null +++ b/quantize_model.py @@ -0,0 +1,47 @@ +#!/usr/bin/env python3 +"""Quantize Ornith-1.0-35B to 4-bit NF4 using bitsandbytes (recommended way).""" + +import torch +from transformers import AutoModelForCausalLM, BitsAndBytesConfig, AutoTokenizer + +def quantize_model(model_path, output_path): + print(f"Quantizing model from: {model_path}") + print("This will use bitsandbytes NF4 quantization with double quantization.\n") + + bnb_config = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.bfloat16, + bnb_4bit_use_double_quant=True, + ) + + print("Loading model with 4-bit quantization (this may take a while)...") + model = AutoModelForCausalLM.from_pretrained( + model_path, + quantization_config=bnb_config, + device_map="auto", + low_cpu_mem_usage=True, + trust_remote_code=True, + ) + + print("\nSaving quantized model...") + model.save_pretrained(output_path) + + # Also save tokenizer if it exists + try: + tokenizer = AutoTokenizer.from_pretrained(model_path) + tokenizer.save_pretrained(output_path) + except: + pass + + print(f"\n✅ Quantized model saved to: {output_path}") + print("You can now load it with: AutoModelForCausalLM.from_pretrained(..., load_in_4bit=True)") + +if __name__ == "__main__": + import argparse + parser = argparse.ArgumentParser() + parser.add_argument("--model-path", type=str, default="/data/models/Ornith-1.0-35B") + parser.add_argument("--output-path", type=str, default="/data/models/Ornith-1.0-35B-4bit-nf4") + args = parser.parse_args() + + quantize_model(args.model_path, args.output_path)