diff --git a/quantize_streaming.py b/quantize_streaming.py new file mode 100644 index 0000000..f38e8b3 --- /dev/null +++ b/quantize_streaming.py @@ -0,0 +1,105 @@ +#!/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()