cmedina/agenx-lora-training
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: load shard to CPU, quantize on GPU

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 22:40:31 -04:00
parent d769d39865
commit 1751d33f69
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
test_quantize_single_shard.py
View File

@@ -26,7 +26,7 @@ def test_single_shard(model_path, output_dir):
shard_file = shards[0] shard_file = shards[0]
print(f"Testing with single shard: {shard_file}") 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 = [ weight_keys = [
k for k, v in state_dict.items() 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 for key in weight_keys[:5]: # Test first 5 layers
try: try:
weight = state_dict[key] weight = state_dict[key].to("cuda:0")
qweight, qstate = quantize_nf4(weight, blocksize=64, compress_statistics=True) qweight, qstate = quantize_nf4(weight, blocksize=64, compress_statistics=True)
state_dict[key] = qweight state_dict[key] = qweight.cpu()
quantized += 1 quantized += 1
print(f"✓ {key}") print(f"✓ {key}")
except Exception as e: except Exception as e: