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

exposed trust_remote_code through the UI

Browse source
This commit is contained in:
samit 2026-03-08 16:28:56 -07:00
commit 6aa50d353f
20 changed files with 94 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

10
studio/backend/core/export/export.py
View file

@ -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

12
studio/backend/core/inference/inference.py
View file

@ -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

2
studio/backend/core/inference/orchestrator.py
View file

@ -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.

1
studio/backend/core/inference/worker.py
View file

@ -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:

19
studio/backend/core/training/trainer.py
View file

@ -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

1
studio/backend/core/training/training.py
View file

@ -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

1
studio/backend/core/training/worker.py
View file

@ -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:

4
studio/backend/models/export.py
View file

@ -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):

4
studio/backend/models/inference.py
View file

@ -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):

4
studio/backend/models/training.py
View file

@ -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")

1
studio/backend/routes/export.py
View file

@ -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:

1
studio/backend/routes/inference.py
View file

@ -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:

1
studio/backend/routes/training.py
View file

@ -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

33
studio/frontend/src/features/chat/chat-settings-sheet.tsx
View file

@ -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({
</CollapsibleSection>
<CollapsibleSection icon={Settings02Icon} label="Settings">
<div className="flex items-center justify-between gap-3 py-1">
<div className="min-w-0">
<div className="text-xs font-medium">Auto title</div>
<div className="text-[11px] text-muted-foreground">
Generate short title after reply.
<div className="flex flex-col gap-3 py-1">
<div className="flex items-center justify-between gap-3">
<div className="min-w-0">
<div className="text-xs font-medium">Auto title</div>
<div className="text-[11px] text-muted-foreground">
Generate short title after reply.
</div>
</div>
<Switch
checked={autoTitle}
onCheckedChange={onAutoTitleChange}
/>
</div>
<div className="flex items-center justify-between gap-3">
<div className="min-w-0">
<div className="text-xs font-medium">Trust remote code</div>
<div className="text-[11px] text-muted-foreground">
Allow models with custom code (e.g. Nemotron). Only enable for repos you trust.
</div>
</div>
<Switch
checked={params.trustRemoteCode ?? false}
onCheckedChange={set("trustRemoteCode")}
/>
</div>
<Switch
checked={autoTitle}
onCheckedChange={onAutoTitleChange}
/>
</div>
</CollapsibleSection>
</div>

6
studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts
View file

@ -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();
}

2
studio/frontend/src/features/chat/types/api.ts
View file

@ -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 {

3
studio/frontend/src/features/chat/types/runtime.ts
View file

@ -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 {

2
studio/frontend/src/features/export/api/export-api.ts
View file

@ -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<ExportOperationResponse> {
const response = await authFetch("/api/export/load-checkpoint", {
method: "POST",

1
studio/frontend/src/features/training/api/mappers.ts
View file

@ -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,

2
studio/frontend/src/features/training/types/api.ts
View file

@ -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;