ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fast LoRA saving

Browse source
This commit is contained in:
Daniel Han-Chen 2024-01-23 03:47:28 +11:00
commit 20d8f223d7
3 changed files with 18 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

13
unsloth/kernels/utils.py
View file

@ -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

8
unsloth/models/llama.py
View file

@ -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

2
unsloth/models/mistral.py
View file

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