feat: true streaming NF4 quantization with safetensors
This commit is contained in:
Christian Medina
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user