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

fix: use BnB quantize_batch to actually quantize weights

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 20:22:39 -04:00
parent ec2b1be864
commit b624e90b02
1 changed files with 15 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
quantize_to_bnb.py
View File

@@ -29,11 +29,20 @@ 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:
# Create new Linear4bit with proper quantization # Quantize weights using BnB
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,
@@ -42,11 +51,11 @@ def quantize_model(model_path, output_path):
quant_type='nf4', quant_type='nf4',
) )
# Copy weights (BnB will quantize during forward) # Set quantized weights
with torch.no_grad(): new_module.weight = nn.Parameter(quantized_weight.view_as(module.weight.data))
new_module.weight.data = module.weight.data.clone() new_module.quant_state = quant_state
if module.bias is not None: if module.bias is not None:
new_module.bias.data = 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('.')