diff --git a/quantize_proper_bnb.py b/quantize_proper_bnb.py new file mode 100644 index 0000000..f343ab3 --- /dev/null +++ b/quantize_proper_bnb.py @@ -0,0 +1,49 @@ +#!/usr/bin/env python3 +"""Quantize Ornith-1.0-35B to 4-bit NF4 (recommended method for 2x RTX 5090)""" + +from transformers import AutoModelForCausalLM, BitsAndBytesConfig, AutoConfig +import torch +import os + +def quantize_model(): + model_path = "/data/models/Ornith-1.0-35B" + output_path = "/data/models/Ornith-1.0-35B-4bit-nf4" + + print(f"Quantizing model: {model_path}") + print("Using 4-bit NF4 with double quantization + aggressive offloading...\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, + ) + + model = AutoModelForCausalLM.from_pretrained( + model_path, + quantization_config=bnb_config, + device_map="auto", + max_memory={ + 0: "26GiB", # Good balance for 5090 (leaves headroom) + 1: "26GiB", + "cpu": "150GiB", # Heavy CPU offloading during quantization + }, + low_cpu_mem_usage=True, + torch_dtype=torch.bfloat16, + trust_remote_code=True, + ) + + print(f"\nSaving quantized model to: {output_path}") + os.makedirs(output_path, exist_ok=True) + model.save_pretrained(output_path) + + # Also save the config explicitly + config = AutoConfig.from_pretrained(model_path, trust_remote_code=True) + config.save_pretrained(output_path) + + print(f"\n✅ Quantization complete!") + print(f" Model saved to: {output_path}") + + +if __name__ == "__main__": + quantize_model()