feat: add streaming quantization script (process one shard at a time)

2026-07-02 21:28:05 -04:00
parent 170243de25
commit 0ea51488a7
1 changed files with 105 additions and 0 deletions
--- a/quantize_streaming.py
+++ 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()