cmedina/agenx-lora-training
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: simple NF4 quantization with device_map=auto (proven method)

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 22:32:38 -04:00
parent bfa569f4e9
commit bee6683396
1 changed files with 17 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
quantize_model.py
View File

@@ -1,12 +1,12 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
"""Quantize Ornith-1.0-35B to 4-bit NF4 using bitsandbytes (recommended way).""" """Simple NF4 quantization using BnB with device_map auto-distribution."""
import torch import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig, AutoTokenizer from transformers import AutoModelForCausalLM, BitsAndBytesConfig, AutoTokenizer
def quantize_model(model_path, output_path): def quantize_model(model_path, output_path):
print(f"Quantizing model from: {model_path}") print(f"Quantizing model from: {model_path}")
print("This will use bitsandbytes NF4 quantization with double quantization.\n")
bnb_config = BitsAndBytesConfig( bnb_config = BitsAndBytesConfig(
load_in_4bit=True, load_in_4bit=True,
@@ -15,27 +15,35 @@ def quantize_model(model_path, output_path):
bnb_4bit_use_double_quant=True, bnb_4bit_use_double_quant=True,
) )
print("Loading model with 4-bit quantization (this may take a while)...") print("Loading model with 4-bit quantization...")
print(" This will distribute across both GPUs automatically\n")
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
model_path, model_path,
quantization_config=bnb_config, quantization_config=bnb_config,
device_map="auto", device_map="auto",
max_memory={0: "28GiB", 1: "28GiB"}, # Leave room for training
low_cpu_mem_usage=True, low_cpu_mem_usage=True,
trust_remote_code=True, trust_remote_code=True,
) )
print("\nSaving quantized model...") print("\n✓ Model loaded and quantized")
print(f" GPU 0: {torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
print(f" GPU 1: {torch.cuda.memory_allocated(1) / 1e9:.2f} GB")
print(f"\nSaving quantized model to: {output_path}")
model.save_pretrained(output_path) model.save_pretrained(output_path)
# Also save tokenizer if it exists # Save tokenizer
try: try:
tokenizer = AutoTokenizer.from_pretrained(model_path) tokenizer = AutoTokenizer.from_pretrained(model_path)
tokenizer.save_pretrained(output_path) tokenizer.save_pretrained(output_path)
print("✓ Tokenizer saved")
except: except:
pass print("⚠ No tokenizer found")
print(f"\n✅ Quantized model saved to: {output_path}") 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__": if __name__ == "__main__":
import argparse import argparse
@@ -43,5 +51,5 @@ if __name__ == "__main__":
parser.add_argument("--model-path", type=str, default="/data/models/Ornith-1.0-35B") 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-nf4") parser.add_argument("--output-path", type=str, default="/data/models/Ornith-1.0-35B-nf4")
args = parser.parse_args() args = parser.parse_args()
quantize_model(args.model_path, args.output_path) quantize_model(args.model_path, args.output_path)