From 1751d33f69c3cace83a2e9288058143482bbd0ed Mon Sep 17 00:00:00 2001 From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com> Date: Thu, 2 Jul 2026 22:40:31 -0400 Subject: [PATCH] fix: load shard to CPU, quantize on GPU --- test_quantize_single_shard.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/test_quantize_single_shard.py b/test_quantize_single_shard.py index c235595..d2487f3 100644 --- a/test_quantize_single_shard.py +++ b/test_quantize_single_shard.py @@ -26,7 +26,7 @@ def test_single_shard(model_path, output_dir): shard_file = shards[0] print(f"Testing with single shard: {shard_file}") - state_dict = load_file(shard_file, device="cuda:0") + state_dict = load_file(shard_file, device="cpu") weight_keys = [ k for k, v in state_dict.items() @@ -40,9 +40,9 @@ def test_single_shard(model_path, output_dir): for key in weight_keys[:5]: # Test first 5 layers try: - weight = state_dict[key] + weight = state_dict[key].to("cuda:0") qweight, qstate = quantize_nf4(weight, blocksize=64, compress_statistics=True) - state_dict[key] = qweight + state_dict[key] = qweight.cpu() quantized += 1 print(f"✓ {key}") except Exception as e: