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()