110 lines
3.6 KiB
Python
110 lines
3.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Properly quantize model to BnB 4-bit using BnB API."""
|
|
|
|
import argparse
|
|
import gc
|
|
import torch
|
|
from pathlib import Path
|
|
from transformers import AutoModelForCausalLM, AutoConfig
|
|
|
|
|
|
def quantize_model(model_path, output_path):
|
|
"""Load bf16 model, properly quantize to BnB 4-bit, save."""
|
|
|
|
print(f"Loading model from: {model_path}")
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
model_path,
|
|
device_map="cpu",
|
|
torch_dtype=torch.bfloat16,
|
|
trust_remote_code=True,
|
|
low_cpu_mem_usage=True,
|
|
)
|
|
print(f"✓ Model loaded to CPU (~70GB bf16)")
|
|
|
|
# Count parameters
|
|
total_params = sum(p.numel() for p in model.parameters())
|
|
print(f" Total parameters: {total_params / 1e9:.2f}B")
|
|
|
|
# Quantize using BnB's actual API
|
|
print("\nQuantizing with BnB 4-bit...")
|
|
from bitsandbytes.nn import Linear4bit
|
|
from torch import nn
|
|
|
|
quantized_count = 0
|
|
for name, module in list(model.named_modules()):
|
|
if isinstance(module, nn.Linear) and 'lm_head' not in name:
|
|
# Create new Linear4bit
|
|
new_module = Linear4bit(
|
|
module.in_features,
|
|
module.out_features,
|
|
bias=module.bias is not None,
|
|
compute_dtype=torch.float16,
|
|
quant_type='nf4',
|
|
)
|
|
|
|
# Copy weights - Linear4bit will quantize on first forward
|
|
with torch.no_grad():
|
|
new_module.weight = nn.Parameter(module.weight.data.clone())
|
|
if module.bias is not None:
|
|
new_module.bias = nn.Parameter(module.bias.data.clone())
|
|
|
|
# Replace in model
|
|
layers = name.split('.')
|
|
parent = model
|
|
for layer in layers[:-1]:
|
|
parent = getattr(parent, layer)
|
|
setattr(parent, layers[-1], new_module)
|
|
quantized_count += 1
|
|
|
|
print(f"✓ Replaced {quantized_count} Linear layers with Linear4bit")
|
|
|
|
# Force quantization by running a dummy forward pass
|
|
print(" Forcing quantization with dummy input...")
|
|
dummy_input = torch.randn(1, 32, model.config.hidden_size)
|
|
with torch.no_grad():
|
|
try:
|
|
model(dummy_input)
|
|
except Exception as e:
|
|
print(f" (Dummy forward may fail, but weights should be quantized)")
|
|
|
|
# Count quantized parameters
|
|
bnb_params = sum(
|
|
1 for p in model.parameters()
|
|
if hasattr(p, 'quant_state') and p.quant_state is not None
|
|
)
|
|
print(f" Quantized modules: {bnb_params}")
|
|
|
|
# Save model config
|
|
print(f"\nSaving to: {output_path}")
|
|
model.config.save_pretrained(output_path)
|
|
|
|
# Save weights
|
|
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()
|