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

TorchAO

Browse source
This commit is contained in:
Daniel Han 2025-10-16 03:55:20 -07:00
commit 92e281f2ee
3 changed files with 38 additions and 20 deletions
Download patch file Download diff file Expand all files Collapse all files

4
unsloth/__init__.py
View file

@ -85,7 +85,7 @@ def get_device_type():
# Check torch.accelerator
if hasattr(torch, "accelerator"):
if not torch.accelerator.is_available():
raise NotImplementedError("Unsloth cannot find any torch accelerators? You need a GPU.")
raise NotImplementedError("Unsloth cannot find any torch accelerator? You need a GPU.")
accelerator = str(torch.accelerator.current_accelerator())
if accelerator in ("cuda", "xpu", "hip"):
raise RuntimeError(
@ -160,6 +160,8 @@ from .import_fixes import ignore_logger_messages
ignore_logger_messages(); del ignore_logger_messages;
from .import_fixes import patch_ipykernel_hf_xet
patch_ipykernel_hf_xet(); del patch_ipykernel_hf_xet;
from .import_fixes import patch_trackio
patch_trackio(); del patch_trackio;
# Torch 2.4 has including_emulation
if DEVICE_TYPE == "cuda":

9
unsloth/import_fixes.py
View file

@ -158,3 +158,12 @@ def patch_ipykernel_hf_xet():
disable_progress_bars()
pass
pass
def patch_trackio():
# Set some environment variables to customize the Trackio dashboard for experiment tracking
# See https://github.com/unslothai/notebooks/pull/110
os.environ["TRACKIO_LOGO_LIGHT_URL"] = "https://raw.githubusercontent.com/unslothai/unsloth/main/images/unsloth%20logo%20black%20text.png"
os.environ["TRACKIO_LOGO_DARK_URL"] = "https://raw.githubusercontent.com/unslothai/unsloth/main/images/unsloth%20logo%20white%20text.png"
os.environ["TRACKIO_PLOT_ORDER"] = "train/loss"
pass
pass

45
unsloth/save.py
View file

@ -13,6 +13,7 @@
# limitations under the License.
from unsloth_zoo.utils import Version
from importlib.metadata import version as importlib_version
from unsloth_zoo.hf_utils import dtype_from_config, HAS_TORCH_DTYPE
from unsloth_zoo.llama_cpp import convert_to_gguf, quantize_gguf, use_local_gguf, install_llama_cpp, check_llama_cpp, _download_convert_hf_to_gguf
from bitsandbytes.nn import Linear4bit as Bnb_Linear4bit
@ -2571,11 +2572,13 @@ def unsloth_save_pretrained_torchao(
del arguments["self"]
del arguments["torchao_config"]
if token is None and push_to_hub: token = get_token()
if not isinstance(self, PeftModelForCausalLM) and not isinstance(self, PeftModel):
self.save_pretrained(save_directory)
tokenizer.save_pretrained(save_directory)
self.save_pretrained(save_directory)
tokenizer.save_pretrained(save_directory)
else:
unsloth_generic_save(**arguments)
unsloth_generic_save(**arguments)
for _ in range(3):
gc.collect()
@ -2583,14 +2586,17 @@ def unsloth_save_pretrained_torchao(
from torchao import quantize_
if torchao_config is None:
from torchao.quantization import Int8DynamicActivationInt8WeightConfig
print("Unsloth: You did not specify a `torchao_config`, so defaulting to `Int8DynamicActivationInt8WeightConfig`")
torchao_config = Int8DynamicActivationInt8WeightConfig()
quantization_config = TorchAoConfig(quant_type = torchao_config)
tokenizer = AutoTokenizer.from_pretrained(arguments["save_directory"])
# TorchAO must only use bfloat16 for loading (float16 fails)
if HAS_TORCH_DTYPE:
kwargs = {"torch_dtype" : "auto"}
kwargs = {"torch_dtype" : torch.bfloat16}
else:
kwargs = {"dtype" : "auto"}
kwargs = {"dtype" : torch.bfloat16}
model = AutoModel.from_pretrained(
arguments["save_directory"],
device_map = "auto",
@ -2600,13 +2606,14 @@ def unsloth_save_pretrained_torchao(
torchao_save_directory = save_directory + "-torchao"
# TorchAO does not support safe_serialization right now for 0.13.0. 0.14.0 yes!
safe_serialization = Version(importlib_version("torchao")) >= Version("0.14.0")
if push_to_hub:
if token is None and push_to_hub: token = get_token()
# torchao does not support safe_serialization right now
model.push_to_hub(torchao_save_directory, safe_serialization = False, token = token)
model.push_to_hub(torchao_save_directory, safe_serialization = safe_serialization, token = token)
tokenizer.push_to_hub(torchao_save_directory, token = token)
else:
model.save_pretrained(torchao_save_directory, safe_serialization=False)
model.save_pretrained(torchao_save_directory, safe_serialization = safe_serialization)
tokenizer.save_pretrained(torchao_save_directory)
pass
for _ in range(3):
@ -2711,20 +2718,20 @@ def patch_saving_functions(model, vision = False):
if not vision:
if hasattr(model, "config"):
# Counteract tokenizers
model.push_to_hub_merged = types.MethodType(unsloth_generic_push_to_hub_merged, model)
model.save_pretrained_merged = types.MethodType(unsloth_generic_save_pretrained_merged, model)
model.push_to_hub_gguf = types.MethodType(unsloth_push_to_hub_gguf, model)
model.save_pretrained_gguf = types.MethodType(unsloth_save_pretrained_gguf, model)
model.save_pretrained_torchao = types.MethodType(unsloth_save_pretrained_torchao, model)
model.push_to_hub_ggml = types.MethodType(unsloth_convert_lora_to_ggml_and_push_to_hub, model)
model.save_pretrained_ggml = types.MethodType(unsloth_convert_lora_to_ggml_and_save_locally, model)
model.push_to_hub_merged = types.MethodType(unsloth_generic_push_to_hub_merged, model)
model.save_pretrained_merged = types.MethodType(unsloth_generic_save_pretrained_merged, model)
model.push_to_hub_gguf = types.MethodType(unsloth_push_to_hub_gguf, model)
model.save_pretrained_gguf = types.MethodType(unsloth_save_pretrained_gguf, model)
model.save_pretrained_torchao = types.MethodType(unsloth_save_pretrained_torchao, model)
model.push_to_hub_ggml = types.MethodType(unsloth_convert_lora_to_ggml_and_push_to_hub, model)
model.save_pretrained_ggml = types.MethodType(unsloth_convert_lora_to_ggml_and_save_locally, model)
pass
else:
# Vision only 1 option
model.push_to_hub_merged = types.MethodType(unsloth_generic_push_to_hub_merged, model)
model.save_pretrained_merged = types.MethodType(unsloth_generic_save_pretrained_merged, model)
model.push_to_hub_gguf = types.MethodType(unsloth_push_to_hub_gguf, model)
model.save_pretrained_gguf = types.MethodType(unsloth_save_pretrained_gguf, model)
model.push_to_hub_merged = types.MethodType(unsloth_generic_push_to_hub_merged, model)
model.save_pretrained_merged = types.MethodType(unsloth_generic_save_pretrained_merged, model)
model.push_to_hub_gguf = types.MethodType(unsloth_push_to_hub_gguf, model)
model.save_pretrained_gguf = types.MethodType(unsloth_save_pretrained_gguf, model)
model.save_pretrained_torchao = types.MethodType(unsloth_save_pretrained_torchao, model)
pass
return model