106 lines
3.6 KiB
Python
106 lines
3.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Streaming quantization: process one shard at a time to avoid OOM."""
|
|
|
|
import argparse
|
|
import gc
|
|
import torch
|
|
from pathlib import Path
|
|
from transformers import AutoConfig, AutoModelForCausalLM
|
|
from bitsandbytes.nn import Linear4bit
|
|
from torch import nn
|
|
|
|
|
|
def streaming_quantize(model_path, output_path):
|
|
"""Quantize model by processing one shard at a time."""
|
|
|
|
print(f"Loading config from: {model_path}")
|
|
config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
|
|
|
|
# Find all safetensors files
|
|
import glob
|
|
safetensors_files = sorted(glob.glob(f"{model_path}/*.safetensors"))
|
|
print(f"Found {len(safetensors_files)} safetensors files")
|
|
|
|
# Create output directory
|
|
output_path = Path(output_path)
|
|
output_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Process each shard
|
|
for shard_idx, shard_file in enumerate(safetensors_files):
|
|
print(f"\n{'='*60}")
|
|
print(f"Processing shard {shard_idx + 1}/{len(safetensors_files)}: {shard_file}")
|
|
print(f"{'='*60}")
|
|
|
|
# Load shard to CPU
|
|
print(" Loading shard to CPU...")
|
|
shard_state_dict = torch.load(shard_file, map_location="cpu", weights_only=True)
|
|
|
|
# Quantize Linear layers in this shard
|
|
print(" Quantizing Linear layers...")
|
|
quantized_keys = 0
|
|
for key, tensor in list(shard_state_dict.items()):
|
|
if 'weight' in key and tensor.dim() == 2:
|
|
# This is a Linear layer weight
|
|
# Create a dummy Linear4bit to get the quantization format
|
|
in_features = tensor.size(1)
|
|
out_features = tensor.size(0)
|
|
|
|
dummy_linear = Linear4bit(
|
|
in_features,
|
|
out_features,
|
|
bias=False,
|
|
compute_dtype=torch.float16,
|
|
quant_type='nf4',
|
|
)
|
|
|
|
# Quantize the weight
|
|
with torch.no_grad():
|
|
dummy_linear.weight = nn.Parameter(tensor.clone())
|
|
# Force quantization by accessing quant_state
|
|
_ = dummy_linear.weight.quant_state
|
|
|
|
# Replace with quantized version
|
|
shard_state_dict[key] = dummy_linear.weight
|
|
quantized_keys += 1
|
|
|
|
print(f" ✓ Quantized {quantized_keys} weights")
|
|
|
|
# Save quantized shard
|
|
shard_name = f"model_shard_{shard_idx:05d}.safetensors"
|
|
shard_path = output_path / shard_name
|
|
print(f" Saving to: {shard_path}")
|
|
torch.save(shard_state_dict, shard_path)
|
|
|
|
# Free memory
|
|
del shard_state_dict
|
|
gc.collect()
|
|
|
|
# Save config
|
|
print(f"\n{'='*60}")
|
|
print("Saving model config...")
|
|
config.save_pretrained(output_path)
|
|
print(f"✓ Model saved to: {output_path}")
|
|
|
|
# Free memory
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
|
print("\n✓ Streaming quantization complete!")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Streaming quantization")
|
|
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-streaming-4bit",
|
|
help="Output path for quantized model")
|
|
args = parser.parse_args()
|
|
|
|
streaming_quantize(args.model_path, args.output_path)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|