From b95b055b4acddfc8132b48500d8d6539e0a15f3a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 14 May 2026 04:33:49 -0700 Subject: [PATCH] studio: comment out training_args.bin torch.load fallback (#5419) 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. --- studio/backend/utils/models/model_config.py | 58 +++++++++++---------- 1 file changed, 30 insertions(+), 28 deletions(-) diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 9eacf893ad..ebf85c5320 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -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