114 lines
3.9 KiB
Python
114 lines
3.9 KiB
Python
#!/usr/bin/env python3
|
|
"""True streaming 4-bit NF4 quantization - one shard at a time."""
|
|
|
|
import argparse
|
|
import gc
|
|
import torch
|
|
from pathlib import Path
|
|
from safetensors.torch import load_file, save_file
|
|
from bitsandbytes.functional import quantize_nf4
|
|
from transformers import AutoConfig
|
|
|
|
|
|
def quantize_weight_nf4(weight: torch.Tensor):
|
|
"""Quantize a single weight tensor to NF4 using bitsandbytes functional API."""
|
|
if weight.dim() != 2:
|
|
return weight, None
|
|
|
|
# quantize_nf4 returns (quantized_tensor, quant_state)
|
|
qweight, quant_state = quantize_nf4(
|
|
weight,
|
|
blocksize=64,
|
|
compress_statistics=True,
|
|
)
|
|
return qweight, quant_state
|
|
|
|
|
|
def streaming_quantize(model_path: str, output_path: str):
|
|
print(f"Streaming NF4 quantization: {model_path}")
|
|
output_path = Path(output_path)
|
|
output_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
|
|
|
|
import glob
|
|
shards = sorted(glob.glob(f"{model_path}/*.safetensors"))
|
|
print(f"Found {len(shards)} shards\n")
|
|
|
|
# Process 4 shards per GPU (8 total) for max parallelism
|
|
shards_per_gpu = 4
|
|
num_gpus = 2
|
|
max_parallel = shards_per_gpu * num_gpus
|
|
|
|
def process_shard(idx, shard_file, gpu_id):
|
|
"""Process a single shard (called per-GPU)."""
|
|
shard_name = f"model-{idx:05d}-of-{len(shards):05d}.safetensors"
|
|
if (output_path / shard_name).exists():
|
|
return
|
|
|
|
print(f"[{idx+1}/{len(shards)}] {Path(shard_file).name} (GPU {gpu_id})")
|
|
|
|
# Load to assigned GPU
|
|
state_dict = load_file(shard_file, device=f"cuda:{gpu_id}")
|
|
|
|
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" Quantizing {len(weight_keys)} tensors...")
|
|
|
|
for key in weight_keys:
|
|
try:
|
|
weight = state_dict[key]
|
|
qweight, qstate = quantize_weight_nf4(weight)
|
|
state_dict[key] = qweight
|
|
del weight, qweight, qstate
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
except Exception as e:
|
|
print(f" Warning: Failed on {key}: {e}")
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
|
# Move to CPU and save
|
|
state_dict = {
|
|
k: v.cpu() if isinstance(v, torch.Tensor) else v
|
|
for k, v in state_dict.items()
|
|
}
|
|
|
|
save_file(state_dict, output_path / shard_name)
|
|
print(f" ✓ Saved {shard_name}")
|
|
|
|
del state_dict
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
|
# Process in batches (4 per GPU = 8 total)
|
|
for i in range(0, len(shards), max_parallel):
|
|
batch = shards[i:i+max_parallel]
|
|
if all((output_path / f"model-{idx:05d}-of-{len(shards):05d}.safetensors").exists() for idx in range(i, i+len(batch))):
|
|
print(f"Shards {i+1}-{i+len(batch)} already saved, skipping")
|
|
continue
|
|
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=max_parallel) as executor:
|
|
futures = []
|
|
for j, shard_file in enumerate(batch):
|
|
idx = i + j
|
|
gpu_id = (j // shards_per_gpu) % num_gpus # Assign to GPU
|
|
futures.append(executor.submit(process_shard, idx, shard_file, gpu_id))
|
|
for future in concurrent.futures.as_completed(futures):
|
|
future.result()
|
|
|
|
config.save_pretrained(output_path)
|
|
print(f"\n✅ Done → {output_path}")
|
|
|
|
|
|
if __name__ == "__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-nf4")
|
|
args = parser.parse_args()
|
|
|
|
streaming_quantize(args.model_path, args.output_path)
|