diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index d78c429e5a..04da0d2db9 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -68,6 +68,7 @@ class UnslothTrainer: self.is_training = False self.should_stop = False self.save_on_stop = True + self.load_in_4bit = True # Track quantization mode for metadata # Model state tracking self.is_vlm = False @@ -114,6 +115,7 @@ class UnslothTrainer: hf_token: Optional[str] = None, is_dataset_multimodal: 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 try: if self.model is not None: del self.model @@ -988,6 +990,7 @@ class UnslothTrainer: # Stopped by user — save model at current checkpoint self.trainer.save_model() self.tokenizer.save_pretrained(output_dir) + self._patch_adapter_config(output_dir) print(f"\nTraining stopped. Model saved to {output_dir}\n") self._update_progress( is_training=False, @@ -1004,6 +1007,7 @@ class UnslothTrainer: # Normal completion self.trainer.save_model() self.tokenizer.save_pretrained(output_dir) + self._patch_adapter_config(output_dir) print(f"\nTraining completed! Model saved to {output_dir}\n") self._update_progress( is_training=False, @@ -1018,6 +1022,36 @@ class UnslothTrainer: finally: self.is_training = False + def _patch_adapter_config(self, output_dir: str) -> None: + """Patch adapter_config.json with unsloth_training_method. + + Values: 'qlora', 'lora', 'FT', 'CPT', 'DPO', 'GRPO', etc. + For LoRA/QLoRA, the distinction comes from load_in_4bit. + """ + config_path = os.path.join(output_dir, "adapter_config.json") + if not os.path.exists(config_path): + logger.info("No adapter_config.json found — skipping training method patch") + return + + try: + with open(config_path, "r") as f: + config = json.load(f) + + # Determine the training method + if self.load_in_4bit: + method = "qlora" + else: + method = "lora" + + config["unsloth_training_method"] = method + logger.info(f"Patching adapter_config.json with unsloth_training_method='{method}'") + + with open(config_path, "w") as f: + json.dump(config, f, indent=2) + + except Exception as e: + logger.warning(f"Failed to patch adapter_config.json: {e}") + def stop_training(self, save: bool = True): """Stop ongoing training""" print(f"\nStopping training (save={save})...") diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index c238b4019c..63fe246008 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -91,11 +91,50 @@ async def load_model(request: LoadRequest): detail=f"Invalid model identifier: {request.model_path}" ) + # Auto-detect quantization for LoRA adapters from adapter_config.json + # The training pipeline patches this file with "unsloth_training_method" + # which is 'qlora' or 'lora'. Only LoRA (16-bit) needs load_in_4bit=False. + load_in_4bit = request.load_in_4bit + if config.is_lora and config.path: + import json + from pathlib import Path + adapter_cfg_path = Path(config.path) / "adapter_config.json" + if adapter_cfg_path.exists(): + try: + with open(adapter_cfg_path) as f: + adapter_cfg = json.load(f) + training_method = adapter_cfg.get("unsloth_training_method") + if training_method == "lora" and load_in_4bit: + logger.info( + f"adapter_config.json says unsloth_training_method='lora' — " + f"setting load_in_4bit=False to match 16-bit training" + ) + load_in_4bit = False + elif training_method == "qlora" and not load_in_4bit: + logger.info( + f"adapter_config.json says unsloth_training_method='qlora' — " + f"setting load_in_4bit=True to match QLoRA training" + ) + load_in_4bit = True + elif training_method: + logger.info(f"Training method: {training_method}, load_in_4bit={load_in_4bit}") + else: + # No unsloth_training_method — fallback to base model name + if config.base_model and "-bnb-4bit" not in config.base_model.lower() and load_in_4bit: + logger.info( + f"No unsloth_training_method in adapter_config.json. " + f"Base model '{config.base_model}' has no -bnb-4bit suffix — " + f"setting load_in_4bit=False" + ) + load_in_4bit = False + except Exception as e: + logger.warning(f"Could not read adapter_config.json: {e}") + # Load the model success = backend.load_model( config=config, max_seq_length=request.max_seq_length, - load_in_4bit=request.load_in_4bit, + load_in_4bit=load_in_4bit, hf_token=request.hf_token, )