fix: replace with proper BnB 4-bit quantization script

quantize_to_bnb.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-"""Quantize bf16 model to BnB 4-bit by replacing Linear layers."""
+"""Properly quantize model to BnB 4-bit using BnB API."""
 
 import argparse
 import gc
@@ -9,7 +9,7 @@ from transformers import AutoModelForCausalLM, AutoConfig
 
 
 def quantize_model(model_path, output_path):
-    """Load bf16 model, quantize to BnB 4-bit, save."""
+    """Load bf16 model, properly quantize to BnB 4-bit, save."""
     
     print(f"Loading model from: {model_path}")
     model = AutoModelForCausalLM.from_pretrained(
@@ -25,31 +25,34 @@ def quantize_model(model_path, output_path):
     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...")
+    # Apply PEFT prepare for k-bit training
+    print("\nApplying PEFT prepare_model_for_kbit_training...")
+    from peft import prepare_model_for_kbit_training
+    model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=False)
+    print("✓ Model prepared for k-bit training")
+    
+    # Quantize using BnB's actual API
+    print("Quantizing 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 4-bit version
+            # Create new Linear4bit with proper quantization
             new_module = Linear4bit(
                 module.in_features,
                 module.out_features,
                 bias=module.bias is not None,
                 compute_dtype=torch.float16,
+                quant_type='nf4',
             )
             
-            # Copy weights
+            # Copy weights (BnB will quantize during forward)
             with torch.no_grad():
-                new_module.weight = nn.Parameter(
-                    module.weight.data.clone()
-                )
+                new_module.weight.data = module.weight.data.clone()
                 if module.bias is not None:
-                    new_module.bias = nn.Parameter(
-                        module.bias.data.clone()
-                    )
+                    new_module.bias.data = module.bias.data.clone()
             
             # Replace in model
             layers = name.split('.')
@@ -61,15 +64,18 @@ def quantize_model(model_path, output_path):
     
     print(f"✓ Quantized {quantized_count} linear layers to 4-bit")
     
-    # Count 4-bit parameters
+    # Count quantized parameters
     bnb_params = sum(
-        p.numel() for p in model.parameters() 
-        if hasattr(p, 'quant_state')
+        1 for p in model.parameters() 
+        if hasattr(p, 'quant_state') and p.quant_state is not None
     )
-    print(f"   4-bit parameters: {bnb_params / 1e9:.2f}B")
+    print(f"   Quantized modules: {bnb_params}")
     
-    # Save model
+    # 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")
     
@@ -79,6 +85,7 @@ def quantize_model(model_path, output_path):
     torch.cuda.empty_cache()
     
     print("\nDone! Model is ready for QLoRA training.")
+    print(f"Save location: {output_path}")
 
 
 def main():