diff --git a/quantize_to_bnb.py b/quantize_to_bnb.py new file mode 100644 index 0000000..500afdf --- /dev/null +++ b/quantize_to_bnb.py @@ -0,0 +1,99 @@ +#!/usr/bin/env python3 +"""Quantize bf16 model to BnB 4-bit by replacing Linear layers.""" + +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, 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") + + # Replace Linear layers with Linear4bit + print("\nReplacing Linear layers with Linear4bit...") + 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 4-bit version + new_module = Linear4bit( + module.in_features, + module.out_features, + bias=module.bias is not None, + compute_dtype=torch.float16, + ) + + # Copy weights + 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"✓ Quantized {quantized_count} linear layers to 4-bit") + + # Count 4-bit parameters + bnb_params = sum( + p.numel() for p in model.parameters() + if hasattr(p, 'quant_state') + ) + print(f" 4-bit parameters: {bnb_params / 1e9:.2f}B") + + # 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.") + + +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()