71 lines
2.1 KiB
Python
71 lines
2.1 KiB
Python
#!/usr/bin/env python3
|
|
"""Quantize bf16 model to BnB 4-bit using transformers' built-in mechanism."""
|
|
|
|
import argparse
|
|
import gc
|
|
import torch
|
|
from pathlib import Path
|
|
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
|
|
|
|
|
|
def quantize_model(model_path, output_path):
|
|
"""Load bf16 model, quantize to BnB 4-bit, save."""
|
|
|
|
print(f"Loading model from: {model_path}")
|
|
|
|
# Load with BnB quantization config and device_map="cpu"
|
|
bnb_config = BitsAndBytesConfig(
|
|
load_in_4bit=True,
|
|
bnb_4bit_quant_type="nf4",
|
|
bnb_4bit_compute_dtype=torch.float16,
|
|
)
|
|
|
|
print("Loading with BnB 4-bit quantization to CPU...")
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
model_path,
|
|
quantization_config=bnb_config,
|
|
device_map="cpu",
|
|
torch_dtype=torch.float16,
|
|
trust_remote_code=True,
|
|
low_cpu_mem_usage=True,
|
|
)
|
|
print("✓ Model loaded with BnB 4-bit to CPU")
|
|
|
|
# Check if model is actually quantized
|
|
bnb_modules = sum(
|
|
1 for m in model.modules()
|
|
if hasattr(m, 'weight') and hasattr(m.weight, 'quant_state')
|
|
)
|
|
print(f" BnB quantized modules: {bnb_modules}")
|
|
|
|
# Save model
|
|
print(f"\nSaving to: {output_path}")
|
|
model.save_pretrained(output_path)
|
|
print("✓ Model saved")
|
|
|
|
# Free memory
|
|
del model
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
|
print("\nDone! Model is ready for QLoRA training.")
|
|
print(f"Save location: {output_path}")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Quantize model to BnB 4-bit")
|
|
parser.add_argument("--model-path", type=str,
|
|
default="/data/models/Ornith-1.0-35B",
|
|
help="Path to bf16 model")
|
|
parser.add_argument("--output-path", type=str,
|
|
default="/data/models/Ornith-1.0-35B-bnb-4bit",
|
|
help="Output path for quantized model")
|
|
args = parser.parse_args()
|
|
|
|
Path(args.output_path).mkdir(parents=True, exist_ok=True)
|
|
quantize_model(args.model_path, args.output_path)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|