chore: clean up temp files

This commit is contained in:
Christian Medina
2026-07-02 20:18:13 -04:00
parent 3716051abc
commit ec2b1be864

View File

@@ -1,106 +0,0 @@
#!/usr/bin/env python3
"""Properly quantize model to BnB 4-bit using BnB API."""
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, properly 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")
# 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 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 (BnB will quantize during forward)
with torch.no_grad():
new_module.weight.data = module.weight.data.clone()
if module.bias is not None:
new_module.bias.data = 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 quantized parameters
bnb_params = sum(
1 for p in model.parameters()
if hasattr(p, 'quant_state') and p.quant_state is not None
)
print(f" Quantized modules: {bnb_params}")
# 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")
# Free memory
del model
gc.collect()
torch.cuda.empty_cache()
print("\nDone! Model is ready for QLoRA training.")
print(f"Save location: {output_path}")
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()