From 959c43d44c9703876a0c108fd2009d1a75c1419c Mon Sep 17 00:00:00 2001 From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com> Date: Thu, 2 Jul 2026 22:05:59 -0400 Subject: [PATCH] feat: true streaming NF4 quantization with safetensors --- quantize_streaming.py | 198 ++++++++++++++++-------------------------- 1 file changed, 77 insertions(+), 121 deletions(-) diff --git a/quantize_streaming.py b/quantize_streaming.py index 4ba362c..29305d6 100644 --- a/quantize_streaming.py +++ b/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): - """Quantize model by processing one shard at a time, using both GPUs.""" - - 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 +def quantize_tensor_to_4bit(tensor: torch.Tensor) -> torch.Tensor: + """Quantize a single 2D weight tensor to 4-bit NF4.""" + in_features = tensor.size(1) + out_features = tensor.size(0) + + # Create a temporary Linear4bit layer + linear_4bit = Linear4bit( + in_features, + out_features, + bias=False, + compute_dtype=torch.bfloat16, + 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 - 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 directly to GPU 0 - print(" Loading shard to GPU 0...") - shard_state_dict = torch.load(shard_file, map_location="cuda:0", weights_only=False) - - # Quantize Linear layers in this shard using both GPUs - print(" Quantizing Linear layers (both GPUs)...") - quantized_keys = 0 - - # Get all weight tensors - weight_keys = [k for k, v in shard_state_dict.items() if 'weight' in k and v.dim() == 2] - - # Distribute between GPUs - gpu0_keys = weight_keys[::2] # Even indices - gpu1_keys = weight_keys[1::2] # Odd indices - - # Quantize on GPU 0 (shard already on GPU 0) - print(" GPU 0: Quantizing...", end=" ") - for key in gpu0_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(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 + + # Load config + config = AutoConfig.from_pretrained(model_path, trust_remote_code=True) + + # Find all shards + import glob + shards = sorted(glob.glob(f"{model_path}/*.safetensors")) + print(f"Found {len(shards)} shards\n") + + for idx, shard_path in enumerate(shards): + print(f"[{idx+1}/{len(shards)}] Processing {Path(shard_path).name}...") + + # Load shard to GPU 0 + state_dict = load_file(shard_path, device="cuda:0") + + # Find all 2D weight tensors (Linear layers) + weight_keys = [k for k, v in state_dict.items() + if "weight" in k and isinstance(v, torch.Tensor) and v.dim() == 2] + + print(f" Found {len(weight_keys)} weight tensors to quantize") + + # Quantize + for key in weight_keys: + try: + state_dict[key] = quantize_tensor_to_4bit(state_dict[key]) + except Exception as e: + print(f" Warning: Failed to quantize {key}: {e}") + + # Move to CPU and save + state_dict = {k: v.cpu() if isinstance(v, torch.Tensor) else v + for k, v in state_dict.items()} + + out_shard = output_path / f"model-{idx:05d}-of-{len(shards):05d}.safetensors" + save_file(state_dict, out_shard) + + print(f" Saved → {out_shard.name}") + + # Cleanup + del state_dict gc.collect() torch.cuda.empty_cache() - - # Save config - print(f"\n{'='*60}") - print("Saving model config...") + + # Save config and index 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") - -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) + # Create model.safetensors.index.json if needed (for multi-shard) + print(f"\n✅ Streaming quantization complete!") + print(f" Output: {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)