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

studio: comment out training_args.bin torch.load fallback (#5419)

Browse source 
torch.load defaults to weights_only=True since torch 2.6, which rejects
the pickled TrainingArguments dataclass that HF Trainer saves to
training_args.bin. Studio ships on torch 2.9 / 2.10 so this fallback
was already failing on every call, getting swallowed by the surrounding
try/except, and falling through to the existing adapter_config.json /
config.json / directory-name paths that already produce the answer.

In get_base_model_from_lora the path is also reachable via the
GET /loras/{lora_path:path}/base-model route on user-supplied paths
(including third-party LoRAs pulled from HF), so "fixing" it with
weights_only=False would re-introduce a pickle deserialization sink
on remote-supplied input.

Comment both blocks out and leave a TODO so the intent is preserved
for whoever wants to re-enable this with proper safe_globals or a
trust check.
This commit is contained in:
Daniel Han 2026-05-14 04:33:49 -07:00 • committed by GitHub
commit b95b055b4a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 30 additions and 28 deletions
Download patch file Download diff file Expand all files Collapse all files

58
studio/backend/utils/models/model_config.py
View file

@ -1696,20 +1696,21 @@ def get_base_model_from_checkpoint(checkpoint_path: str) -> Optional[str]:
)
return base_model
training_args_path = checkpoint_path_obj / "training_args.bin"
if training_args_path.exists():
try:
import torch
training_args = torch.load(training_args_path)
if hasattr(training_args, "model_name_or_path"):
base_model = training_args.model_name_or_path
logger.info(
"Detected base model from training_args.bin: %s", base_model
)
return base_model
except Exception as e:
logger.warning(f"Could not load training_args.bin: {e}")
# TODO: torch.load default weights_only=True (torch >= 2.6) rejects pickled TrainingArguments; re-enable via safe_globals or weights_only=False once threat model allows.
# training_args_path = checkpoint_path_obj / "training_args.bin"
# if training_args_path.exists():
# try:
# import torch
#
# training_args = torch.load(training_args_path)
# if hasattr(training_args, "model_name_or_path"):
# base_model = training_args.model_name_or_path
# logger.info(
# "Detected base model from training_args.bin: %s", base_model
# )
# return base_model
# except Exception as e:
# logger.warning(f"Could not load training_args.bin: {e}")
dir_name = checkpoint_path_obj.name
if dir_name.startswith("unsloth_"):
@ -1757,20 +1758,21 @@ def get_base_model_from_lora(lora_path: str) -> Optional[str]:
return base_model
# Fallback: try training_args.bin (requires torch)
training_args_path = lora_path_obj / "training_args.bin"
if training_args_path.exists():
try:
import torch
training_args = torch.load(training_args_path)
if hasattr(training_args, "model_name_or_path"):
base_model = training_args.model_name_or_path
logger.info(
f"Detected base model from training_args.bin: {base_model}"
)
return base_model
except Exception as e:
logger.warning(f"Could not load training_args.bin: {e}")
# TODO: torch.load default weights_only=True (torch >= 2.6) rejects pickled TrainingArguments; also an RCE sink for third-party LoRAs via this route, re-enable behind a trust check if needed.
# training_args_path = lora_path_obj / "training_args.bin"
# if training_args_path.exists():
# try:
# import torch
#
# training_args = torch.load(training_args_path)
# if hasattr(training_args, "model_name_or_path"):
# base_model = training_args.model_name_or_path
# logger.info(
# f"Detected base model from training_args.bin: {base_model}"
# )
# return base_model
# except Exception as e:
# logger.warning(f"Could not load training_args.bin: {e}")
# Last resort: parse from directory name
# Format: unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit_timestamp