From 0666589acec7aeff93da8aa7684adf4ab2c75e2c Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 22 Jan 2024 04:28:24 +1100 Subject: [PATCH] Update save.py --- unsloth/save.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/unsloth/save.py b/unsloth/save.py index 64342d91de..519d8e77d8 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -77,10 +77,14 @@ def _merge_lora(layer, name): W, quant_state, A, B, s = get_lora_parameters(layer) dtype = quant_state.dtype if type(quant_state) is not list else quant_state[2] W = fast_dequantize(W, quant_state).to(torch.float32).t() - sAB = (A.t().to(torch.float32) @ (s * B.t().to(torch.float32))) - W += sAB - if not torch.isfinite(W).all(): - raise ValueError(f"Unsloth: Merge failed.\n{name} has some elements = infinity.") + + if A is not None: + sAB = (A.t().to(torch.float32) @ (s * B.t().to(torch.float32))) + W += sAB + if not torch.isfinite(W).all(): + raise ValueError(f"Unsloth: Merge failed.\n{name} has some elements = infinity.") + pass + W = W.t().to(dtype) else: W = layer.weight @@ -156,6 +160,7 @@ def unsloth_save_model( pass if save_method == "merged_4bit": + print("Unsloth: Merging 4bit and LoRA weights to 4bit...") print("This might take 5 minutes...") model = model.merge_and_unload()