Enable native ModelOpt quantization support (3/3) (#10154)

Signed-off-by: Zhiyu Cheng <zhiyuc@nvidia.com>

python/sglang/srt/model_loader/loader.py

@@ -538,12 +538,21 @@ class DefaultModelLoader(BaseModelLoader):
             **model_kwargs,
             trust_remote_code=True,
         )
-        rank0_log(f"ModelOpt quantization requested: {model_config.modelopt_quant}")
+        # Handle both legacy modelopt_quant and unified quantization flags
+        if hasattr(model_config, "modelopt_quant") and model_config.modelopt_quant:
+            # Legacy approach
+            quant_choice_str = model_config.modelopt_quant
+            rank0_log(f"ModelOpt quantization requested (legacy): {quant_choice_str}")
+        else:
+            # Unified approach - extract quantization type
+            quant_choice_str = model_config._get_modelopt_quant_type()
+            rank0_log(
+                f"ModelOpt quantization requested (unified): {model_config.quantization} -> {quant_choice_str}"
+            )
 
-        quant_choice_str = model_config.modelopt_quant
         if not isinstance(quant_choice_str, str):
             raise TypeError(
-                f"modelopt_quant must be a string preset key (e.g., 'fp8'), "
+                f"Quantization type must be a string (e.g., 'fp8'), "
                 f"got {type(quant_choice_str)}"
             )
 
@@ -1764,6 +1773,7 @@ class ModelOptModelLoader(DefaultModelLoader):
         quant_cfg,
         quantized_ckpt_restore_path: str | None = None,
         quantized_ckpt_save_path: str | None = None,
+        export_path: str | None = None,
     ) -> None:
         """
         Set up ModelOpt quantization for the given model.
@@ -1774,6 +1784,7 @@ class ModelOptModelLoader(DefaultModelLoader):
             quant_cfg: The quantization configuration
             quantized_ckpt_restore_path: Path to restore quantized checkpoint from
             quantized_ckpt_save_path: Path to save quantized checkpoint to
+            export_path: Path to export the quantized model in HuggingFace format
 
         Raises:
             ImportError: If ModelOpt is not available
@@ -1798,6 +1809,9 @@ class ModelOptModelLoader(DefaultModelLoader):
                 rank0_log(
                     f"Restored quantized model from {quantized_ckpt_restore_path}"
                 )
+
+                # Export model if path provided (even when restoring from checkpoint)
+                self._maybe_export_modelopt(model, export_path)
                 return
             except Exception as e:
                 logger.warning(
@@ -1844,9 +1858,75 @@ class ModelOptModelLoader(DefaultModelLoader):
                         f"Failed to save quantized checkpoint to {quantized_ckpt_save_path}: {e}"
                     )
 
+            # Export model if path provided
+            self._maybe_export_modelopt(model, export_path)
+
         except Exception as e:
             raise Exception(f"Failed to set up ModelOpt quantization: {e}") from e
 
+    def _maybe_export_modelopt(self, model, export_path: str | None) -> None:
+        """Export model to HuggingFace format if export_path is provided."""
+        if export_path:
+            try:
+                # Get the original model path from the model config
+                original_model_path = getattr(self, "_original_model_path", None)
+                self._export_modelopt_checkpoint(
+                    model, export_path, original_model_path
+                )
+                rank0_log(
+                    f"Quantized model exported to HuggingFace format at {export_path}"
+                )
+            except Exception as e:
+                rank0_log(
+                    f"Warning: Failed to export quantized model to {export_path}: {e}"
+                )
+
+    def _export_modelopt_checkpoint(
+        self,
+        model,
+        export_path: str,
+        model_path: str = None,
+        trust_remote_code: bool = True,
+    ) -> None:
+        """
+        Export the quantized model to HuggingFace format using ModelOpt export API.
+
+        Args:
+            model: The quantized model to export
+            export_path: Directory path to export the model to
+            model_path: Path to the original model (for tokenizer export)
+            trust_remote_code: Whether to trust remote code for tokenizer loading
+
+        Raises:
+            ImportError: If ModelOpt export functionality is not available
+            Exception: If export fails
+        """
+        try:
+            from modelopt.torch.export import export_hf_checkpoint
+            from transformers import AutoTokenizer
+        except ImportError as e:
+            raise ImportError(
+                "ModelOpt export functionality is not available. "
+                "Please ensure you have the latest version of modelopt installed."
+            ) from e
+
+        # Create export directory if it doesn't exist
+        os.makedirs(export_path, exist_ok=True)
+
+        # Export the quantized model
+        export_hf_checkpoint(model, export_dir=export_path)
+
+        # Export the tokenizer if model_path is provided
+        if model_path:
+            try:
+                tokenizer = AutoTokenizer.from_pretrained(
+                    model_path, trust_remote_code=trust_remote_code
+                )
+                tokenizer.save_pretrained(export_path)
+                rank0_log(f"Tokenizer exported to {export_path}")
+            except Exception as e:
+                rank0_log(f"Warning: Failed to export tokenizer: {e}")
+
     def load_model(
         self,
         *,
@@ -1856,28 +1936,52 @@ class ModelOptModelLoader(DefaultModelLoader):
 
         logger.info("ModelOptModelLoader: Loading base model...")
 
-        # Use shared method from parent class to load base model
+        # Store the original model path for tokenizer export
+        self._original_model_path = model_config.model_path
+
+        # Check if model is already quantized
+        if model_config._is_already_quantized():
+            logger.info("Model is already quantized, loading directly...")
+            # Use default loading for pre-quantized models
+            return super().load_model(
+                model_config=model_config, device_config=device_config
+            )
+
+        # TODO: Quantize-and-serve mode has been disabled at the ModelConfig level
+        # All quantization now uses the standard workflow (quantize + export/save)
+        logger.info("Standard quantization mode: Will quantize and export/save")
+        return self._standard_quantization_workflow(model_config, device_config)
+
+    def _standard_quantization_workflow(
+        self, model_config: ModelConfig, device_config: DeviceConfig
+    ) -> nn.Module:
+        """Standard quantization workflow: quantize, save checkpoint, export, then return model."""
+        # Use shared method from parent class to load base model for quantization
         model = self._load_modelopt_base_model(model_config)
 
-        # Import ModelOpt modules (already done in _load_modelopt_base_model, but needed here for quantization)
+        # Import ModelOpt modules
         try:
             import modelopt.torch.quantization as mtq
         except ImportError:
             logger.error(
                 "NVIDIA Model Optimizer (modelopt) library not found. "
-                "Please install it to use 'modelopt_quant' feature."
+                "Please install it to use ModelOpt quantization."
             )
             raise
 
-        quant_choice_str = model_config.modelopt_quant
+        # Handle both old modelopt_quant and new unified quantization flags
+        if hasattr(model_config, "modelopt_quant") and model_config.modelopt_quant:
+            # Legacy modelopt_quant flag
+            quant_choice_str = model_config.modelopt_quant
+        else:
+            # Unified quantization flag - extract the type (fp8/fp4)
+            quant_choice_str = model_config._get_modelopt_quant_type()
 
         quant_cfg_name = QUANT_CFG_CHOICES.get(quant_choice_str)
         if not quant_cfg_name:
             raise ValueError(
-                f"Invalid modelopt_quant choice: '{quant_choice_str}'. "
-                f"Available choices in QUANT_CFG_CHOICES: {list(QUANT_CFG_CHOICES.keys())}. "
-                "Ensure QUANT_CFG_CHOICES is correctly defined with mappings to "
-                "attribute names of config objects in modelopt.torch.quantization."
+                f"Invalid quantization choice: '{quant_choice_str}'. "
+                f"Available choices: {list(QUANT_CFG_CHOICES.keys())}"
             )
 
         try:
@@ -1885,20 +1989,27 @@ class ModelOptModelLoader(DefaultModelLoader):
             quant_cfg = getattr(mtq, quant_cfg_name)
         except AttributeError:
             raise AttributeError(
-                f"ModelOpt quantization config attribute '{quant_cfg_name}' "
-                f"(from choice '{quant_choice_str}') not found in modelopt.torch.quantization module. "
-                "Please verify QUANT_CFG_CHOICES and the ModelOpt library."
+                f"ModelOpt quantization config '{quant_cfg_name}' not found. "
+                "Please verify the ModelOpt library installation."
             )
 
         logger.info(
-            f"Quantizing model with ModelOpt using config attribute: mtq.{quant_cfg_name}"
+            f"Quantizing model with ModelOpt using config: mtq.{quant_cfg_name}"
         )
 
-        quantized_ckpt_restore_path = model_config.modelopt_checkpoint_restore_path
-        quantized_ckpt_save_path = model_config.modelopt_checkpoint_save_path
+        # Get ModelOpt configuration from LoadConfig
+        modelopt_config = self.load_config.modelopt_config
+        quantized_ckpt_restore_path = (
+            modelopt_config.checkpoint_restore_path if modelopt_config else None
+        )
+        quantized_ckpt_save_path = (
+            modelopt_config.checkpoint_save_path if modelopt_config else None
+        )
+        export_path = modelopt_config.export_path if modelopt_config else None
         tokenizer = AutoTokenizer.from_pretrained(
             model_config.model_path, use_fast=True
         )
+
         try:
             self._setup_modelopt_quantization(
                 model,
@@ -1906,6 +2017,7 @@ class ModelOptModelLoader(DefaultModelLoader):
                 quant_cfg,
                 quantized_ckpt_restore_path=quantized_ckpt_restore_path,
                 quantized_ckpt_save_path=quantized_ckpt_save_path,
+                export_path=export_path,
             )
         except Exception as e:
             logger.warning(f"ModelOpt quantization failed: {e}")
@@ -1919,12 +2031,27 @@ def get_model_loader(
 ) -> BaseModelLoader:
     """Get a model loader based on the load format."""
 
+    if model_config and (
+        (hasattr(model_config, "modelopt_quant") and model_config.modelopt_quant)
+        or model_config.quantization in ["modelopt_fp8", "modelopt_fp4", "modelopt"]
+    ):
+        logger.info("Using ModelOptModelLoader due to ModelOpt quantization config.")
+        return ModelOptModelLoader(load_config)
+
+    # Use ModelOptModelLoader for unified quantization flags
     if (
         model_config
-        and hasattr(model_config, "modelopt_quant")
-        and model_config.modelopt_quant
+        and hasattr(model_config, "quantization")
+        and model_config.quantization in ["modelopt_fp8", "modelopt_fp4"]
     ):
-        logger.info("Using ModelOptModelLoader due to 'modelopt_quant' config.")
+        if model_config._is_already_quantized():
+            logger.info(
+                f"Using ModelOptModelLoader for pre-quantized model: {model_config.quantization}"
+            )
+        else:
+            logger.info(
+                f"Using ModelOptModelLoader for quantization: {model_config.quantization}"
+            )
         return ModelOptModelLoader(load_config)
 
     if isinstance(load_config.load_format, type):