cmedina/agenx-lora-training
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: force quantization with dummy forward pass

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 20:25:26 -04:00
parent b624e90b02
commit 11f9c3a56c
1 changed files with 16 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
quantize_to_bnb.py
View File

@@ -29,20 +29,11 @@ def quantize_model(model_path, output_path):
print("\nQuantizing with BnB 4-bit...") print("\nQuantizing with BnB 4-bit...")
from bitsandbytes.nn import Linear4bit from bitsandbytes.nn import Linear4bit
from torch import nn from torch import nn
from bitsandbytes import quantize_batch
quantized_count = 0 quantized_count = 0
for name, module in list(model.named_modules()): for name, module in list(model.named_modules()):
if isinstance(module, nn.Linear) and 'lm_head' not in name: if isinstance(module, nn.Linear) and 'lm_head' not in name:
# Quantize weights using BnB # Create new Linear4bit
weight_2d = module.weight.data.view(-1, module.weight.data.size(-1))
quantized_weight, quant_state = quantize_batch(
weight_2d.to(torch.float16),
blocksize=64,
quant_type='nf4',
)
# Create new Linear4bit with quantized weights
new_module = Linear4bit( new_module = Linear4bit(
module.in_features, module.in_features,
module.out_features, module.out_features,
@@ -51,11 +42,11 @@ def quantize_model(model_path, output_path):
quant_type='nf4', quant_type='nf4',
) )
# Set quantized weights # Copy weights - Linear4bit will quantize on first forward
new_module.weight = nn.Parameter(quantized_weight.view_as(module.weight.data)) with torch.no_grad():
new_module.quant_state = quant_state new_module.weight = nn.Parameter(module.weight.data.clone())
if module.bias is not None: if module.bias is not None:
new_module.bias = nn.Parameter(module.bias.data.clone()) new_module.bias = nn.Parameter(module.bias.data.clone())
# Replace in model # Replace in model
layers = name.split('.') layers = name.split('.')
@@ -65,7 +56,16 @@ def quantize_model(model_path, output_path):
setattr(parent, layers[-1], new_module) setattr(parent, layers[-1], new_module)
quantized_count += 1 quantized_count += 1
print(f"✓ Quantized {quantized_count} linear layers to 4-bit") 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 # Count quantized parameters
bnb_params = sum( bnb_params = sum(