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

fix: don't save quant_state objects (BnB rebuilds on load)

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 22:48:09 -04:00
parent 8c13be4d1f
commit 18ef0f73ef
1 changed files with 3 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
quantize_streaming.py
View File

@@ -47,16 +47,12 @@ def streaming_quantize(model_path: str, output_path: str):
print(f" Quantizing {len(weight_keys)} tensors...") print(f" Quantizing {len(weight_keys)} tensors...")
quant_states = {} # We need to save these too for loading later
for key in weight_keys: for key in weight_keys:
try: try:
weight = state_dict[key] weight = state_dict[key]
qweight, qstate = quantize_weight_nf4(weight) qweight, qstate = quantize_weight_nf4(weight)
state_dict[key] = qweight state_dict[key] = qweight
if qstate is not None: del weight, qweight, qstate
quant_states[f"{key}.quant_state"] = qstate
del weight, qweight
gc.collect() gc.collect()
torch.cuda.empty_cache() torch.cuda.empty_cache()
except Exception as e: except Exception as e:
@@ -64,19 +60,14 @@ def streaming_quantize(model_path: str, output_path: str):
gc.collect() gc.collect()
torch.cuda.empty_cache() torch.cuda.empty_cache()
# Move everything to CPU # Move to CPU and save
state_dict = { state_dict = {
k: v.cpu() if isinstance(v, torch.Tensor) else v k: v.cpu() if isinstance(v, torch.Tensor) else v
for k, v in state_dict.items() 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" shard_name = f"model-{idx:05d}-of-{len(shards):05d}.safetensors"
save_file({**state_dict, **quant_states}, output_path / shard_name) save_file(state_dict, output_path / shard_name)
print(f" Saved {shard_name}") print(f" Saved {shard_name}")