feat: add BnB 4-bit quantization script

2026-07-02 20:04:44 -04:00
parent 3c31e4dfb0
commit 76ff09b140
1 changed files with 99 additions and 0 deletions
--- a/quantize_to_bnb.py
+++ 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()