fix: free GPU memory after each layer quantization
This commit is contained in:
@@ -51,12 +51,18 @@ def streaming_quantize(model_path: str, output_path: str):
|
||||
|
||||
for key in weight_keys:
|
||||
try:
|
||||
qweight, qstate = quantize_weight_nf4(state_dict[key])
|
||||
weight = state_dict[key]
|
||||
qweight, qstate = quantize_weight_nf4(weight)
|
||||
state_dict[key] = qweight
|
||||
if qstate is not None:
|
||||
quant_states[f"{key}.quant_state"] = qstate
|
||||
del weight, qweight
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
except Exception as e:
|
||||
print(f" Warning: Failed on {key}: {e}")
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# Move everything to CPU
|
||||
state_dict = {
|
||||
|
||||
Reference in New Issue
Block a user