diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index a00be89942..7c693a6c4e 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -181,9 +181,16 @@ def fast_linear_forward(proj, X, temp_lora = None, out = None): W, W_quant, lora_A, lora_B, lora_S = get_lora_parameters(proj) out = fast_gemv(X, W, W_quant, out = out) if lora_A is not None: - dtype = X.dtype - temp_lora = torch.matmul(X, lora_A.to(dtype).t(), out = temp_lora) - out.addmv_(lora_B.to(dtype), temp_lora, alpha = lora_S) + + # Save LoRAs for inference to stop data movement costs + if not hasattr(lora_A, "_fast_lora"): + dtype = X.dtype + lora_A._fast_lora = lora_A.to(dtype).t() + lora_B._fast_lora = lora_B.to(dtype) + pass + + temp_lora = torch.matmul(X, lora_A._fast_lora, out = temp_lora) + out.addmv_(lora_B._fast_lora, temp_lora, alpha = lora_S) pass return out pass diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 18947744c5..40074ad043 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -724,7 +724,7 @@ class FastLlamaModel: f" \\\ /| GPU: {gpu_stats.name}. Max memory: {max_memory} GB. Platform = {platform_system}.\n"\ f"O^O/ \_/ \\ Pytorch: {torch.__version__}. CUDA = {gpu_stats.major}.{gpu_stats.minor}. CUDA Toolkit = {torch.version.cuda}.\n"\ f"\ / Bfloat16 = {str(SUPPORTS_BFLOAT16).upper()}. Xformers = {xformers_version}. FA = {HAS_FLASH_ATTENTION}.\n"\ - f' "-____-" Apache 2 free license - http://github.com/unslothai/unsloth' + f' "-____-" Apache 2 free license: http://github.com/unslothai/unsloth' logger.warning_once(statistics) FastLlamaModel.pre_patch() @@ -1152,6 +1152,12 @@ class FastLlamaModel: internal_model.gradient_checkpointing = use_gradient_checkpointing internal_model.training = True + # Delete all fast inference loras + for param in model.parameters(): + if hasattr(param, "_fast_lora"): + del param._fast_lora + pass + while hasattr(internal_model, "model"): internal_model = internal_model.model internal_model.gradient_checkpointing = use_gradient_checkpointing diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index 55a27ed048..bde68a0a51 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -280,7 +280,7 @@ class FastMistralModel(FastLlamaModel): f" \\\ /| GPU: {gpu_stats.name}. Max memory: {max_memory} GB. Platform = {platform_system}.\n"\ f"O^O/ \_/ \\ Pytorch: {torch.__version__}. CUDA = {gpu_stats.major}.{gpu_stats.minor}. CUDA Toolkit = {torch.version.cuda}.\n"\ f"\ / Bfloat16 = {str(SUPPORTS_BFLOAT16).upper()}. Xformers = {xformers_version}. FA = {HAS_FLASH_ATTENTION}.\n"\ - f' "-____-" Apache 2 free license - http://github.com/unslothai/unsloth' + f' "-____-" Apache 2 free license: http://github.com/unslothai/unsloth' logger.warning_once(statistics) FastMistralModel.pre_patch()