92 lines
2.9 KiB
Python
92 lines
2.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")
|
|
|
|
for idx, shard_file in enumerate(shards):
|
|
print(f"[{idx+1}/{len(shards)}] {Path(shard_file).name}")
|
|
|
|
state_dict = load_file(shard_file, device="cuda:0")
|
|
|
|
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...")
|
|
|
|
quant_states = {} # We need to save these too for loading later
|
|
|
|
for key in weight_keys:
|
|
try:
|
|
qweight, qstate = quantize_weight_nf4(state_dict[key])
|
|
state_dict[key] = qweight
|
|
if qstate is not None:
|
|
quant_states[f"{key}.quant_state"] = qstate
|
|
except Exception as e:
|
|
print(f" Warning: Failed on {key}: {e}")
|
|
|
|
# Move everything to CPU
|
|
state_dict = {
|
|
k: v.cpu() if isinstance(v, torch.Tensor) else v
|
|
for k, v in state_dict.items()
|
|
}
|
|
quant_states = {
|
|
k: v.cpu() if isinstance(v, torch.Tensor) else v
|
|
for k, v in quant_states.items()
|
|
}
|
|
|
|
# Save shard + quant states
|
|
shard_name = f"model-{idx:05d}-of-{len(shards):05d}.safetensors"
|
|
save_file({**state_dict, **quant_states}, output_path / shard_name)
|
|
|
|
print(f" Saved {shard_name}")
|
|
|
|
del state_dict, quant_states
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
|
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)
|