Patch adapter_config.json with unsloth_training_method and auto-detect load_in_4bit for LoRA inference

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})...")

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