feat: true streaming NF4 quantization with safetensors

quantize_streaming.py

@@ -1,140 +1,96 @@
 #!/usr/bin/env python3
-"""Streaming quantization: process one shard at a time to avoid OOM."""
+"""True streaming 4-bit NF4 quantization - one shard at a time."""
 
 import argparse
 import gc
 import torch
 from pathlib import Path
-from transformers import AutoConfig, AutoModelForCausalLM
+from safetensors.torch import load_file, save_file
 from bitsandbytes.nn import Linear4bit
 from torch import nn
+from transformers import AutoConfig
 
 
-def streaming_quantize(model_path, output_path):
+def quantize_tensor_to_4bit(tensor: torch.Tensor) -> torch.Tensor:
-    """Quantize model by processing one shard at a time, using both GPUs."""
+    """Quantize a single 2D weight tensor to 4-bit NF4."""
-    
+    in_features = tensor.size(1)
-    print(f"Loading config from: {model_path}")
+    out_features = tensor.size(0)
-    config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
+
-    
+    # Create a temporary Linear4bit layer
-    # Find all safetensors files
+    linear_4bit = Linear4bit(
-    import glob
+        in_features,
-    safetensors_files = sorted(glob.glob(f"{model_path}/*.safetensors"))
+        out_features,
-    print(f"Found {len(safetensors_files)} safetensors files")
+        bias=False,
-    
+        compute_dtype=torch.bfloat16,
-    # Create output directory
+        quant_type="nf4",
+    ).to(tensor.device)
+
+    with torch.no_grad():
+        linear_4bit.weight = nn.Parameter(tensor.clone())
+        # Force quantization to happen
+        _ = linear_4bit.weight.quant_state
+
+    return linear_4bit.weight
+
+
+def streaming_quantize(model_path: str, output_path: str):
+    print(f"Streaming quantization of: {model_path}")
     output_path = Path(output_path)
     output_path.mkdir(parents=True, exist_ok=True)
-    
+
-    # Process each shard
+    # Load config
-    for shard_idx, shard_file in enumerate(safetensors_files):
+    config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
-        print(f"\n{'='*60}")
+
-        print(f"Processing shard {shard_idx + 1}/{len(safetensors_files)}: {shard_file}")
+    # Find all shards
-        print(f"{'='*60}")
+    import glob
-        
+    shards = sorted(glob.glob(f"{model_path}/*.safetensors"))
-        # Load shard directly to GPU 0
+    print(f"Found {len(shards)} shards\n")
-        print("   Loading shard to GPU 0...")
+
-        shard_state_dict = torch.load(shard_file, map_location="cuda:0", weights_only=False)
+    for idx, shard_path in enumerate(shards):
-        
+        print(f"[{idx+1}/{len(shards)}] Processing {Path(shard_path).name}...")
-        # Quantize Linear layers in this shard using both GPUs
+
-        print("   Quantizing Linear layers (both GPUs)...")
+        # Load shard to GPU 0
-        quantized_keys = 0
+        state_dict = load_file(shard_path, device="cuda:0")
-        
+
-        # Get all weight tensors
+        # Find all 2D weight tensors (Linear layers)
-        weight_keys = [k for k, v in shard_state_dict.items() if 'weight' in k and v.dim() == 2]
+        weight_keys = [k for k, v in state_dict.items() 
-        
+                       if "weight" in k and isinstance(v, torch.Tensor) and v.dim() == 2]
-        # Distribute between GPUs
+
-        gpu0_keys = weight_keys[::2]  # Even indices
+        print(f"   Found {len(weight_keys)} weight tensors to quantize")
-        gpu1_keys = weight_keys[1::2]  # Odd indices
+
-        
+        # Quantize
-        # Quantize on GPU 0 (shard already on GPU 0)
+        for key in weight_keys:
-        print("      GPU 0: Quantizing...", end=" ")
+            try:
-        for key in gpu0_keys:
+                state_dict[key] = quantize_tensor_to_4bit(state_dict[key])
-            tensor = shard_state_dict[key]
+            except Exception as e:
-            in_features = tensor.size(1)
+                print(f"   Warning: Failed to quantize {key}: {e}")
-            out_features = tensor.size(0)
+
-            
+        # Move to CPU and save
-            dummy_linear = Linear4bit(
+        state_dict = {k: v.cpu() if isinstance(v, torch.Tensor) else v 
-                in_features,
+                      for k, v in state_dict.items()}
-                out_features,
+
-                bias=False,
+        out_shard = output_path / f"model-{idx:05d}-of-{len(shards):05d}.safetensors"
-                compute_dtype=torch.float16,
+        save_file(state_dict, out_shard)
-                quant_type='nf4',
+
-            )
+        print(f"   Saved → {out_shard.name}")
-            
+
-            with torch.no_grad():
+        # Cleanup
-                dummy_linear.weight = nn.Parameter(tensor.clone())
+        del state_dict
-                _ = dummy_linear.weight.quant_state
-            
-            shard_state_dict[key] = dummy_linear.weight
-            quantized_keys += 1
-        print(f"✓ {len(gpu0_keys)} layers")
-        
-        # Move shard to GPU 1 for second half
-        print("      Moving to GPU 1...")
-        shard_state_dict = {k: v.to("cuda:1") if isinstance(v, torch.Tensor) else v for k, v in shard_state_dict.items()}
-        
-        # Quantize on GPU 1
-        print("      GPU 1: Quantizing...", end=" ")
-        for key in gpu1_keys:
-            tensor = shard_state_dict[key]
-            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',
-            )
-            
-            with torch.no_grad():
-                dummy_linear.weight = nn.Parameter(tensor.clone())
-                _ = dummy_linear.weight.quant_state
-            
-            shard_state_dict[key] = dummy_linear.weight
-            quantized_keys += 1
-        print(f"✓ {len(gpu1_keys)} layers")
-        
-        print(f"   ✓ Total: {quantized_keys} layers quantized")
-        
-        # 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()
         torch.cuda.empty_cache()
-    
+
-    # Save config
+    # Save config and index
-    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!")
-    print(f"   Used both GPUs in parallel for faster quantization")
 
-
+    # Create model.safetensors.index.json if needed (for multi-shard)
-def main():
+    print(f"\n✅ Streaming quantization complete!")
-    parser = argparse.ArgumentParser(description="Streaming quantization")
+    print(f"   Output: {output_path}")
-    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()
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-path", type=str, default="/data/models/Ornith-1.0-35B")
+    parser.add_argument("--output-path", type=str, default="/data/models/Ornith-1.0-35B-4bit-streaming")
+    args = parser.parse_args()
+
+    streaming_quantize(args.model_path, args.output_path)