From 6aa50d353f12e4317676c3fa7c083800c14239f5 Mon Sep 17 00:00:00 2001 From: samit Date: Sun, 8 Mar 2026 16:28:56 -0700 Subject: [PATCH] exposed trust_remote_code through the UI --- studio/backend/core/export/export.py | 10 +++++- studio/backend/core/inference/inference.py | 12 ++++++- studio/backend/core/inference/orchestrator.py | 2 ++ studio/backend/core/inference/worker.py | 1 + studio/backend/core/training/trainer.py | 19 +++++++++-- studio/backend/core/training/training.py | 1 + studio/backend/core/training/worker.py | 1 + studio/backend/models/export.py | 4 +++ studio/backend/models/inference.py | 4 +++ studio/backend/models/training.py | 4 +++ studio/backend/routes/export.py | 1 + studio/backend/routes/inference.py | 1 + studio/backend/routes/training.py | 1 + .../src/features/chat/chat-settings-sheet.tsx | 33 ++++++++++++++----- .../chat/hooks/use-chat-model-runtime.ts | 6 ++-- .../frontend/src/features/chat/types/api.ts | 2 ++ .../src/features/chat/types/runtime.ts | 3 ++ .../src/features/export/api/export-api.ts | 2 ++ .../src/features/training/api/mappers.ts | 1 + .../src/features/training/types/api.ts | 2 ++ 20 files changed, 94 insertions(+), 16 deletions(-) diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index 9e1c5cff84..bd14fb412e 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -138,7 +138,8 @@ class ExportBackend: def load_checkpoint(self, checkpoint_path: str, max_seq_length: int = 2048, - load_in_4bit: bool = True) -> Tuple[bool, str]: + load_in_4bit: bool = True, + trust_remote_code: bool = False) -> Tuple[bool, str]: """ Load a checkpoint for export. @@ -178,6 +179,7 @@ class ExportBackend: dtype=None, auto_model=CsmForConditionalGeneration, load_in_4bit=False, + trust_remote_code=trust_remote_code, ) elif self._audio_type == 'whisper': @@ -189,6 +191,7 @@ class ExportBackend: dtype=None, load_in_4bit=False, auto_model=WhisperForConditionalGeneration, + trust_remote_code=trust_remote_code, ) elif self._audio_type == 'snac': @@ -198,6 +201,7 @@ class ExportBackend: max_seq_length=max_seq_length, dtype=None, load_in_4bit=load_in_4bit, + trust_remote_code=trust_remote_code, ) elif self._audio_type == 'bicodec': @@ -208,6 +212,7 @@ class ExportBackend: max_seq_length=max_seq_length, dtype=torch.float32, load_in_4bit=False, + trust_remote_code=trust_remote_code, ) elif self._audio_type == 'dac': @@ -217,6 +222,7 @@ class ExportBackend: model_name=checkpoint_path, max_seq_length=max_seq_length, load_in_4bit=False, + trust_remote_code=trust_remote_code, ) elif self.is_vision: @@ -226,6 +232,7 @@ class ExportBackend: max_seq_length=max_seq_length, dtype=None, load_in_4bit=load_in_4bit, + trust_remote_code=trust_remote_code, ) tokenizer = processor # For vision models, processor acts as tokenizer @@ -236,6 +243,7 @@ class ExportBackend: max_seq_length=max_seq_length, dtype=None, load_in_4bit=load_in_4bit, + trust_remote_code=trust_remote_code, ) # Check if PEFT model diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 780a399637..6ced3fc2ad 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -63,7 +63,8 @@ class InferenceBackend: max_seq_length: int = 2048, dtype = None, load_in_4bit: bool = True, - hf_token: Optional[str] = None) -> bool: + hf_token: Optional[str] = None, + trust_remote_code: bool = False) -> bool: """ Load any model: base, LoRA adapter, text, or vision. """ @@ -110,6 +111,7 @@ class InferenceBackend: auto_model=CsmForConditionalGeneration, load_in_4bit=False, token=hf_token if hf_token and hf_token.strip() else None, + trust_remote_code=trust_remote_code, ) FastModel.for_inference(model) self.models[model_name]["model"] = model @@ -139,6 +141,7 @@ class InferenceBackend: dtype=torch.float32, load_in_4bit=False, token=hf_token if hf_token and hf_token.strip() else None, + trust_remote_code=trust_remote_code, ) else: # Base model: download full HF repo, then load from /LLM subfolder @@ -155,6 +158,7 @@ class InferenceBackend: dtype=torch.float32, load_in_4bit=False, token=hf_token if hf_token and hf_token.strip() else None, + trust_remote_code=trust_remote_code, ) FastModel.for_inference(model) @@ -169,6 +173,7 @@ class InferenceBackend: max_seq_length=max_seq_length, load_in_4bit=False, token=hf_token if hf_token and hf_token.strip() else None, + trust_remote_code=trust_remote_code, ) FastModel.for_inference(model) self.models[model_name]["model"] = model @@ -184,6 +189,7 @@ class InferenceBackend: whisper_task="transcribe", load_in_4bit=False, token=hf_token if hf_token and hf_token.strip() else None, + trust_remote_code=trust_remote_code, ) FastModel.for_inference(model) model.eval() @@ -209,6 +215,7 @@ class InferenceBackend: max_seq_length=max_seq_length, load_in_4bit=False, token=hf_token if hf_token and hf_token.strip() else None, + trust_remote_code=trust_remote_code, ) FastLanguageModel.for_inference(model) self.models[model_name]["model"] = model @@ -240,6 +247,7 @@ class InferenceBackend: dtype=dtype, load_in_4bit=load_in_4bit, token=hf_token if hf_token and hf_token.strip() else None, + trust_remote_code=trust_remote_code, ) # Apply inference optimization @@ -270,6 +278,7 @@ class InferenceBackend: processor = AutoProcessor.from_pretrained( processor_source, token=hf_token if hf_token and hf_token.strip() else None, + trust_remote_code=trust_remote_code, ) logger.info(f"Loaded {type(processor).__name__} from {processor_source}") @@ -285,6 +294,7 @@ class InferenceBackend: dtype=dtype, load_in_4bit=load_in_4bit, token=hf_token if hf_token and hf_token.strip() else None, + trust_remote_code=trust_remote_code, ) # Apply inference optimization diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index b08ec5bb70..6910b3dc19 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -258,6 +258,7 @@ class InferenceOrchestrator: dtype=None, load_in_4bit: bool = True, hf_token: Optional[str] = None, + trust_remote_code: bool = False, ) -> bool: """Load a model for inference. @@ -282,6 +283,7 @@ class InferenceOrchestrator: "load_in_4bit": load_in_4bit, "hf_token": hf_token or "", "gguf_variant": getattr(config, "gguf_variant", None), + "trust_remote_code": trust_remote_code, } # Always kill existing subprocess and spawn fresh. diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 3e766b3b01..e59372c216 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -155,6 +155,7 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: max_seq_length=config.get("max_seq_length", 2048), load_in_4bit=load_in_4bit, hf_token=hf_token, + trust_remote_code=config.get("trust_remote_code", False), ) if success: diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index d50431b035..40d653b69a 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -325,9 +325,11 @@ class UnslothTrainer: load_in_4bit: bool = True, hf_token: Optional[str] = None, is_dataset_image: bool = False, - is_dataset_audio: bool = False) -> bool: + is_dataset_audio: bool = False, + trust_remote_code: bool = False) -> bool: """Load model for training (supports both text and vision models)""" self.load_in_4bit = load_in_4bit # Store for training_meta.json + self.trust_remote_code = trust_remote_code # For AutoProcessor etc. used during training try: if self.model is not None: del self.model @@ -446,6 +448,7 @@ class UnslothTrainer: auto_model=CsmForConditionalGeneration, load_in_4bit=False, token=hf_token, + trust_remote_code=trust_remote_code, ) logger.info("Loaded CSM audio model") @@ -461,6 +464,7 @@ class UnslothTrainer: whisper_language="English", whisper_task="transcribe", token=hf_token, + trust_remote_code=trust_remote_code, ) # Configure generation settings (notebook lines 100-105) self.model.generation_config.language = "<|en|>" @@ -477,6 +481,7 @@ class UnslothTrainer: dtype=None, load_in_4bit=load_in_4bit, token=hf_token, + trust_remote_code=trust_remote_code, ) logger.info(f"Loaded {self._audio_type} audio model (FastLanguageModel)") @@ -510,6 +515,7 @@ class UnslothTrainer: dtype=torch.float32, # Spark-TTS requires float32 load_in_4bit=False, token=hf_token, + trust_remote_code=trust_remote_code, ) logger.info("Loaded Spark-TTS (bicodec) model") @@ -521,6 +527,7 @@ class UnslothTrainer: max_seq_length=max_seq_length, load_in_4bit=False, token=hf_token, + trust_remote_code=trust_remote_code, ) logger.info("Loaded OuteTTS (dac) model (FastModel)") @@ -534,6 +541,7 @@ class UnslothTrainer: dtype=None, load_in_4bit=load_in_4bit, token=hf_token, + trust_remote_code=trust_remote_code, ) logger.info("Loaded audio VLM model (FastModel)") @@ -545,6 +553,7 @@ class UnslothTrainer: dtype=None, # Auto-detect load_in_4bit=load_in_4bit, token=hf_token, + trust_remote_code=trust_remote_code, ) logger.info("Loaded vision model") @@ -564,6 +573,7 @@ class UnslothTrainer: dtype=None, # Auto-detect load_in_4bit=load_in_4bit, token=hf_token, + trust_remote_code=trust_remote_code, ) logger.info("Loaded text model") @@ -584,7 +594,7 @@ class UnslothTrainer: self._source_code_retried = True print(f"\n'could not get source code' — retrying once...\n") return self.load_model(model_name, max_seq_length, load_in_4bit, hf_token, - is_dataset_image, is_dataset_audio) + is_dataset_image, is_dataset_audio, trust_remote_code) error_msg = str(e) error_lower = error_msg.lower() if any(k in error_lower for k in ("gated repo", "access to it at", "401", "403", "unauthorized", "forbidden")): @@ -983,7 +993,10 @@ class UnslothTrainer: from datasets import Audio import torch - processor = AutoProcessor.from_pretrained(self.model_name) + processor = AutoProcessor.from_pretrained( + self.model_name, + trust_remote_code=getattr(self, "trust_remote_code", False), + ) # Strip pad_to_multiple_of from tokenizer init_kwargs — fine-tuned models # (e.g. keanteng/sesame-csm-elise) save it in tokenizer_config.json, and diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 5323e73ae1..cc0124ae83 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -176,6 +176,7 @@ class TrainingBackend: "wandb_project": kwargs.get("wandb_project", "unsloth-training"), "enable_tensorboard": kwargs.get("enable_tensorboard", False), "tensorboard_dir": kwargs.get("tensorboard_dir", "runs"), + "trust_remote_code": kwargs.get("trust_remote_code", False), } # Derive load_in_4bit from training_type diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 4c82347936..8eee8293e1 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -192,6 +192,7 @@ def run_training_process( hf_token=hf_token, is_dataset_image=config.get("is_dataset_image", False), is_dataset_audio=config.get("is_dataset_audio", False), + trust_remote_code=config.get("trust_remote_code", False), ) if not success or trainer.should_stop: if trainer.should_stop: diff --git a/studio/backend/models/export.py b/studio/backend/models/export.py index a0526e80db..df17596185 100644 --- a/studio/backend/models/export.py +++ b/studio/backend/models/export.py @@ -19,6 +19,10 @@ class LoadCheckpointRequest(BaseModel): True, description="Whether to load the model in 4-bit quantization", ) + trust_remote_code: bool = Field( + False, + description="Allow loading models with custom code. Only enable for checkpoints/base models you trust.", + ) class ExportStatusResponse(BaseModel): diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 0eb7a7edaf..ab5c63d8c9 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -18,6 +18,10 @@ class LoadRequest(BaseModel): load_in_4bit: bool = Field(True, description="Load model in 4-bit quantization") is_lora: bool = Field(False, description="Whether this is a LoRA adapter") gguf_variant: Optional[str] = Field(None, description="GGUF quantization variant (e.g. 'Q4_K_M')") + trust_remote_code: bool = Field( + False, + description="Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.", + ) class UnloadRequest(BaseModel): diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 5cc8141bac..c03d7b04bf 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -13,6 +13,10 @@ class TrainingStartRequest(BaseModel): hf_token: Optional[str] = Field(None, description="HuggingFace token") load_in_4bit: bool = Field(True, description="Load model in 4-bit quantization") max_seq_length: int = Field(2048, description="Maximum sequence length") + trust_remote_code: bool = Field( + False, + description="Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.", + ) # Dataset parameters hf_dataset: Optional[str] = Field(None, description="HuggingFace dataset identifier") diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index 7f18f50eaa..bf41cba424 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -102,6 +102,7 @@ async def load_checkpoint( checkpoint_path=request.checkpoint_path, max_seq_length=request.max_seq_length, load_in_4bit=request.load_in_4bit, + trust_remote_code=request.trust_remote_code, ) if not success: diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 52e5bfad72..b2f7a802a4 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -226,6 +226,7 @@ async def load_model( max_seq_length=request.max_seq_length, load_in_4bit=load_in_4bit, hf_token=request.hf_token, + trust_remote_code=request.trust_remote_code, ) if not success: diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index d6726d69df..b6e6d7e111 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -191,6 +191,7 @@ async def start_training( "wandb_project": request.wandb_project or "", "enable_tensorboard": request.enable_tensorboard, "tensorboard_dir": request.tensorboard_dir or "", + "trust_remote_code": request.trust_remote_code, } # Free GPU memory: shut down any running inference/export subprocesses diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index e70b24fb54..2ffa5a5370 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -170,6 +170,7 @@ export function ChatSettingsPanel({ ...p.params, systemPrompt: params.systemPrompt, checkpoint: params.checkpoint, + trustRemoteCode: params.trustRemoteCode, }); setActivePreset(name); } @@ -330,17 +331,31 @@ export function ChatSettingsPanel({ -
-
-
Auto title
-
- Generate short title after reply. +
+
+
+
Auto title
+
+ Generate short title after reply. +
+ +
+
+
+
Trust remote code
+
+ Allow models with custom code (e.g. Nemotron). Only enable for repos you trust. +
+
+
-
diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index eb59d5a23d..827da9d68b 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -206,6 +206,7 @@ export function useChatModelRuntime() { await unloadModel({ model_path: currentCheckpoint }); } + const currentParams = useChatRuntimeStore.getState().params; const loadResponse = await loadModel({ model_path: modelId, hf_token: null, @@ -213,10 +214,11 @@ export function useChatModelRuntime() { load_in_4bit: true, is_lora: isLora, gguf_variant: ggufVariant ?? null, + trust_remote_code: currentParams.trustRemoteCode ?? false, }); - const currentParams = useChatRuntimeStore.getState().params; - setParams(mergeRecommendedInference(currentParams, loadResponse, modelId)); + const paramsAfterLoad = useChatRuntimeStore.getState().params; + setParams(mergeRecommendedInference(paramsAfterLoad, loadResponse, modelId)); await refresh(); } diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index 4a2ed9a3f1..9999e0ce70 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -34,6 +34,8 @@ export interface LoadModelRequest { load_in_4bit: boolean; is_lora: boolean; gguf_variant?: string | null; + /** Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust. */ + trust_remote_code?: boolean; } export interface GgufVariantDetail { diff --git a/studio/frontend/src/features/chat/types/runtime.ts b/studio/frontend/src/features/chat/types/runtime.ts index a544c35f64..a1c98ada7d 100644 --- a/studio/frontend/src/features/chat/types/runtime.ts +++ b/studio/frontend/src/features/chat/types/runtime.ts @@ -7,6 +7,8 @@ export interface InferenceParams { maxTokens: number; systemPrompt: string; checkpoint: string; + /** Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust. */ + trustRemoteCode?: boolean; } export const DEFAULT_INFERENCE_PARAMS: InferenceParams = { @@ -18,6 +20,7 @@ export const DEFAULT_INFERENCE_PARAMS: InferenceParams = { maxTokens: 4092, systemPrompt: "", checkpoint: "", + trustRemoteCode: false, }; export interface ChatModelSummary { diff --git a/studio/frontend/src/features/export/api/export-api.ts b/studio/frontend/src/features/export/api/export-api.ts index cbaa4a8145..1eac52f216 100644 --- a/studio/frontend/src/features/export/api/export-api.ts +++ b/studio/frontend/src/features/export/api/export-api.ts @@ -50,6 +50,8 @@ export async function loadCheckpoint(params: { checkpoint_path: string; max_seq_length?: number; load_in_4bit?: boolean; + /** Allow loading models with custom code. Only enable for checkpoints you trust. */ + trust_remote_code?: boolean; }): Promise { const response = await authFetch("/api/export/load-checkpoint", { method: "POST", diff --git a/studio/frontend/src/features/training/api/mappers.ts b/studio/frontend/src/features/training/api/mappers.ts index ca4339838c..572a650b08 100644 --- a/studio/frontend/src/features/training/api/mappers.ts +++ b/studio/frontend/src/features/training/api/mappers.ts @@ -32,6 +32,7 @@ export function buildTrainingStartPayload( hf_token: config.hfToken.trim() || null, load_in_4bit: adapterMethod ? isQlorMethod : false, max_seq_length: config.contextLength, + trust_remote_code: false, hf_dataset: hfDataset, subset: hfDataset ? config.datasetSubset : null, train_split: hfDataset ? config.datasetSplit : null, diff --git a/studio/frontend/src/features/training/types/api.ts b/studio/frontend/src/features/training/types/api.ts index 0dcb18b1bf..2e99930a2c 100644 --- a/studio/frontend/src/features/training/types/api.ts +++ b/studio/frontend/src/features/training/types/api.ts @@ -4,6 +4,8 @@ export interface TrainingStartRequest { hf_token: string | null; load_in_4bit: boolean; max_seq_length: number; + /** Allow loading models with custom code. Only enable for repos you trust. */ + trust_remote_code?: boolean; hf_dataset: string | null; subset: string | null; train_split: string | null;