67 lines
2.0 KiB
Python
67 lines
2.0 KiB
Python
#!/usr/bin/env python3
|
|
"""Test quantization on single shard."""
|
|
|
|
import gc
|
|
import torch
|
|
from pathlib import Path
|
|
from safetensors.torch import load_file
|
|
from bitsandbytes.functional import quantize_nf4
|
|
|
|
|
|
def test_one_shard(model_path, output_dir):
|
|
output_dir = Path(output_dir)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
import glob
|
|
shards = sorted(glob.glob(f"{model_path}/*.safetensors"))
|
|
|
|
if len(shards) == 0:
|
|
print("No shards found!")
|
|
return
|
|
|
|
shard_file = shards[0] # First shard only
|
|
print(f"Testing shard: {shard_file}")
|
|
|
|
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"Found {len(weight_keys)} weight tensors\n")
|
|
|
|
quant_states = {}
|
|
quantized = 0
|
|
failed = 0
|
|
|
|
for key in weight_keys:
|
|
try:
|
|
weight = state_dict[key]
|
|
qweight, qstate = quantize_nf4(weight, blocksize=64, compress_statistics=True)
|
|
state_dict[key] = qweight
|
|
if qstate is not None:
|
|
quant_states[f"{key}.quant_state"] = qstate
|
|
quantized += 1
|
|
print(f"✓ {key} -> {tuple(qweight.shape)}")
|
|
del weight, qweight, qstate
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
except Exception as e:
|
|
failed += 1
|
|
print(f"✗ {key}: {e}")
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
|
print(f"\nResults: {quantized} quantized, {failed} failed, {len(quant_states)} quant states")
|
|
|
|
# Save test output as model.safetensors
|
|
state_dict_cpu = {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in state_dict.items()}
|
|
test_output = output_dir / "model.safetensors"
|
|
torch.save({**state_dict_cpu, **quant_states}, test_output)
|
|
print(f"Saved to: {test_output}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_one_shard("/data/models/Ornith-1.0-35B", "/data/models/test_quantize")
|