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

Update save.py

Browse source
This commit is contained in:
Daniel Han-Chen 2024-02-07 03:45:07 +11:00
commit 2b346dc498
1 changed files with 5 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

12
unsloth/save.py
View file

@ -374,11 +374,11 @@ def unsloth_save_model(
# Save to GPU memory
state_dict[name] = W
# [TODO] Saving to RAM seems to leak memory???
elif (max_ram - W.nbytes) > 0:
# Save to CPU memory
logger.warning_once(f"We will save to RAM and not VRAM now.")
state_dict[name] = W.to("cpu", non_blocking = True, copy = True)
max_ram = max(max_ram - W.nbytes, 0)
# elif (max_ram - W.nbytes) > 0:
# # Save to CPU memory
# logger.warning_once(f"We will save to RAM and not VRAM now.")
# state_dict[name] = W.to("cpu", non_blocking = True, copy = True)
# max_ram = max(max_ram - W.nbytes, 0)
else:
# Save to Disk
logger.warning_once(f"We will save to Disk and not RAM now.")
@ -886,7 +886,6 @@ def unsloth_save_pretrained_gguf(
makefile = install_llama_cpp_make_non_blocking()
new_save_directory = unsloth_save_model(**arguments)
python_install.wait()
makefile.wait()
else:
try:
new_save_directory = unsloth_save_model(**arguments)
@ -899,7 +898,6 @@ def unsloth_save_pretrained_gguf(
makefile = install_llama_cpp_make_non_blocking()
new_save_directory = unsloth_save_model(**arguments)
python_install.wait()
makefile.wait()
pass
pass