Use moe_context in pipeline and by default and add tests

Author: sairampillai

src/llmcompressor/args/dataset_arguments.py

@@ -7,6 +7,7 @@
 HuggingFace datasets, custom JSON/CSV files, and DVC-managed datasets.
 """
 
+import warnings
 from dataclasses import dataclass, field
 from typing import Any, Callable, Dict, List, Optional, Union
 
@@ -126,16 +127,6 @@ class DatasetArguments(CustomDatasetArguments):
         default=512,
         metadata={"help": "Number of samples to use for one-shot calibration"},
     )
-    calibrate_moe_context: bool = field(
-        default=False,
-        metadata={
-            "help": "If during calibration, the MoE context should be enabled "
-            "for the given model. This usually involves updating all MoE modules "
-            "in the model for the duration of calibration. See moe_context under "
-            "modeling/prepare.py for a list of supported MoEs and their updated "
-            "module definitions"
-        },
-    )
     shuffle_calibration_samples: Optional[bool] = field(
         default=True,
         metadata={
@@ -181,6 +172,17 @@ class DatasetArguments(CustomDatasetArguments):
             ),
         },
     )
+    calibrate_moe_context: Optional[bool] = field(
+        default=None,
+        metadata={
+            "help": (
+                "DEPRECATED: This parameter is deprecated and will be \
+                    removed in a future version. "
+                "MoE calibration context is now handled automatically by the pipeline. "
+                "This parameter is ignored and will not affect the calibration process."
+            ),
+        },
+    )
     # --- pipeline arguments --- #
     pipeline: Optional[str] = field(
         default="independent",
@@ -229,3 +231,16 @@ class DatasetArguments(CustomDatasetArguments):
 
     def is_dataset_provided(self) -> bool:
         return self.dataset is not None or self.dataset_path is not None
+
+    def __post_init__(self):
+        """Post-initialization hook to issue deprecation warnings."""
+        if self.calibrate_moe_context is not None:
+            warnings.warn(
+                "The 'calibrate_moe_context' parameter is deprecated\
+                     and will be removed in a future version. "
+                "MoE calibration context is now handled automatically by the pipeline. "
+                "This parameter is ignored and will not affect\
+                     the calibration process.",
+                DeprecationWarning,
+                stacklevel=2,
+            )

src/llmcompressor/modeling/moe_context.py

@@ -8,13 +8,72 @@
 
 import contextlib
 from abc import ABC, abstractmethod
-from typing import Dict, TypeVar, Union
+from dataclasses import dataclass
+from enum import Enum
+from typing import Callable, Dict, Optional, TypeVar, Union
 
+import tqdm
+from compressed_tensors.utils import replace_module
 from transformers import PreTrainedModel
 
+from llmcompressor.utils.helpers import patch_attr
+
 T = TypeVar("T", bound="MoECalibrationContext")
 
 
+class MoECalibrationType(Enum):
+    """Enumeration of supported MoE calibration types."""
+
+    PERMANENT = "permanent"
+    CONTEXTUAL = "contextual"
+
+
+@dataclass
+class MoEModelConfig:
+    """
+    Configuration for MoE model calibration.
+
+    This dataclass defines the parameters needed to configure MoE calibration
+    for a specific model architecture. It follows the same pattern used by
+    other model configuration systems in the project (e.g., SmoothQuant, AWQ).
+
+    Attributes:
+        calibration_type: Type of calibration - MoECalibrationType.PERMANENT or
+            MoECalibrationType.CONTEXTUAL
+        target_class_name: The class name of the MoE module to replace
+        replace_function: Function that creates the replacement module
+        target_attribute: For contextual calibration, the attribute to replace
+        description: Optional description of the model configuration
+    """
+
+    calibration_type: MoECalibrationType
+    target_class_name: str
+    replace_function: Callable
+    target_attribute: Optional[str] = None
+    description: Optional[str] = None
+
+    def __post_init__(self):
+        """Validate configuration after initialization."""
+        if (
+            self.calibration_type == MoECalibrationType.CONTEXTUAL
+            and self.target_attribute is None
+        ):
+            raise ValueError("target_attribute is required for contextual calibration")
+
+        if (
+            self.calibration_type == MoECalibrationType.PERMANENT
+            and self.target_attribute is not None
+        ):
+            raise ValueError(
+                "target_attribute should not be set for permanent calibration"
+            )
+
+
+# Registry of MoE model configurations
+# Add new MoE models here following the same pattern as MAPPINGS_REGISTRY
+MOE_MODEL_REGISTRY: Dict[str, MoEModelConfig] = {}
+
+
 class MoECalibrationContext(ABC):
     """
     Abstract base class for MoE calibration.
@@ -60,13 +119,14 @@ def apply(self, model: PreTrainedModel, calibrate_all_experts: bool = True) -> N
         """Apply MoE calibration modifications using context managers."""
         if self._stack is None:
             self._stack = contextlib.ExitStack()
+            self._stack.__enter__()
 
         self.update_function(model, self._stack, calibrate_all_experts)
 
     def restore(self, model: PreTrainedModel) -> None:
         """Restore the model by exiting the context stack."""
         if self._stack is not None:
-            self._stack.close()
+            self._stack.__exit__(None, None, None)
             self._stack = None
 
 
@@ -134,22 +194,108 @@ def list_supported_models() -> list:
     return list(_MOE_CONTEXTS.keys())
 
 
-# Convenience function for backward compatibility
-def create_context_manager_context(model_class_name: str, update_function):
+# Generic factory functions for creating MoE updaters
+def create_permanent_moe_updater(target_class_name: str, replace_function: Callable):
     """
-    Create a context manager-based MoE calibration.
-    :param model_class_name: The class name of the model
-    :param update_function: Function that applies the MoE modifications
-    :return: A ContextualMoECalibration instance
+    Create a permanent MoE updater function for the given target class.
+
+    Args:
+        target_class_name: The class name to look for in the model
+        replace_function: Function that creates the replacement module
+
+    Returns:
+        A function that can be used with PermanentMoECalibration
     """
-    return ContextualMoECalibration(model_class_name, update_function)
 
+    def update_function(model: PreTrainedModel, calibrate_all_experts: bool):
+        """Update MoE modules for calibration."""
+        for name, module in tqdm.tqdm(list(model.named_modules())):
+            if module.__class__.__name__ == target_class_name:
+                new_module = replace_function(
+                    config=model.config,
+                    module=module,
+                    calibrate_all_experts=calibrate_all_experts,
+                )
+                replace_module(model, name, new_module)
+
+    return update_function
 
-def create_permanent_context(model_class_name: str, replacement_function):
+
+def create_contextual_moe_updater(
+    target_class_name: str, target_attr: str, replace_function: Callable
+):
     """
-    Create a permanent MoE calibration.
-    :param model_class_name: The class name of the model
-    :param replacement_function: Function that permanently replaces MoE modules
-    :return: A PermanentMoECalibration instance
+    Create a contextual MoE updater function for the given target class and attribute.
+
+    Args:
+        target_class_name: The class name to look for in the model
+        target_attr: The attribute name to replace within the target class
+        replace_function: Function that creates the replacement module
+
+    Returns:
+        A function that can be used with ContextualMoECalibration
+    """
+
+    def update_function(
+        model: PreTrainedModel, stack: contextlib.ExitStack, calibrate_all_experts: bool
+    ):
+        """Update MoE modules for calibration using context managers."""
+        for module in model.modules():
+            if module.__class__.__name__ == target_class_name:
+                stack.enter_context(
+                    patch_attr(
+                        module,
+                        target_attr,
+                        replace_function(
+                            config=model.config,
+                            module=getattr(module, target_attr),
+                            calibrate_all_experts=calibrate_all_experts,
+                        ),
+                    )
+                )
+
+    return update_function
+
+
+def register_moe_model(model_class_name: str, config: MoEModelConfig):
     """
-    return PermanentMoECalibration(model_class_name, replacement_function)
+    Register a MoE model with its configuration.
+
+    Args:
+        model_class_name: The model class name
+        config: MoEModelConfig dataclass instance with calibration parameters
+    """
+    if config.calibration_type == MoECalibrationType.PERMANENT:
+        updater = create_permanent_moe_updater(
+            config.target_class_name, config.replace_function
+        )
+        context = PermanentMoECalibration(config.target_class_name, updater)
+    elif config.calibration_type == MoECalibrationType.CONTEXTUAL:
+        updater = create_contextual_moe_updater(
+            config.target_class_name, config.target_attribute, config.replace_function
+        )
+        context = ContextualMoECalibration(model_class_name, updater)
+    else:
+        raise ValueError(f"Unknown MoE type: {config.calibration_type}")
+
+    register_moe_context(model_class_name, context)
+
+
+def register_moe_model_from_dict(model_class_name: str, config_dict: dict):
+    """
+    Register a MoE model from a dictionary configuration (backward compatibility).
+
+    Args:
+        model_class_name: The model class name
+        config_dict: Dictionary with calibration parameters
+    """
+    # Convert string calibration_type to enum
+    if "calibration_type" in config_dict and isinstance(
+        config_dict["calibration_type"], str
+    ):
+        config_dict["calibration_type"] = MoECalibrationType(
+            config_dict["calibration_type"]
+        )
+
+    config = MoEModelConfig(**config_dict)
+    register_moe_model(model_class_name, config)

src/llmcompressor/modeling/prepare.py

@@ -1,3 +1,4 @@
+import contextlib
 import warnings
 
 import tqdm
@@ -6,10 +7,15 @@
 
 from llmcompressor.modeling.deepseek_v3 import replace as replace_deepseekv3
 from llmcompressor.modeling.llama4 import replace as replace_llama4
+from llmcompressor.modeling.moe_context import (
+    MoECalibrationType,
+    MoEModelConfig,
+    get_moe_context,
+    register_moe_model,
+)
 from llmcompressor.modeling.qwen3_moe import replace as replace_Qwen3MoE
-from llmcompressor.utils.helpers import patch_attr
 
-__all__ = ["replace_modules_for_calibration"]
+__all__ = ["moe_calibration_context"]
 
 # ---------------------- module replacements; permanent -------------------------
 replacements = {
@@ -24,8 +30,8 @@ def replace_modules_for_calibration(
 ) -> PreTrainedModel:
     # This function is deprecated. Use moe_calibration_context instead.
     warnings.warn(
-        "replace_modules_for_calibration is deprecated. \
-    Use moe_calibration_context instead.",
+        "replace_modules_for_calibration is deprecated. "
+        "Use moe_calibration_context instead.",
         DeprecationWarning,
         stacklevel=2,
     )
@@ -45,37 +51,67 @@ def replace_modules_for_calibration(
 
 # ------------------- module replacements; during calibration --------------------
 
-
-def update_qwen3_moe(model, stack, calibrate_all_experts):
-    for module in model.modules():
-        cls_name = module.__class__.__name__
-        if cls_name == "Qwen3MoeDecoderLayer":
-            # Optionally update the model.config to pass in other arguments
-            stack.enter_context(
-                patch_attr(
-                    module,
-                    "mlp",
-                    replace_Qwen3MoE(
-                        config=model.config,
-                        module=module.mlp,
-                        calibrate_all_experts=calibrate_all_experts,
-                    ),
-                )
-            )
+# MoE model configurations - centralized registry
+# Adding a new MoE model is now as simple as adding an entry here!
+# This follows the same pattern as MAPPINGS_REGISTRY in SmoothQuant and AWQ
+MOE_EXPERTS_REPLACEMENTS = {
+    "Qwen3MoeForCausalLM": MoEModelConfig(
+        calibration_type=MoECalibrationType.CONTEXTUAL,
+        target_class_name="Qwen3MoeDecoderLayer",
+        target_attribute="mlp",
+        replace_function=replace_Qwen3MoE,
+        description="Qwen3 MoE model with contextual calibration for MLP layers",
+    ),
+    "DeepseekV3ForCausalLM": MoEModelConfig(
+        calibration_type=MoECalibrationType.PERMANENT,
+        target_class_name="DeepseekV3MoE",
+        replace_function=replace_deepseekv3,
+        description="DeepSeek V3 MoE model with permanent calibration",
+    ),
+    "Llama4ForConditionalGeneration": MoEModelConfig(
+        calibration_type=MoECalibrationType.PERMANENT,
+        target_class_name="Llama4TextMoe",
+        replace_function=replace_llama4,
+        description=(
+            "Llama4 MoE model with permanent calibration for vLLM compatibility"
+        ),
+    ),
+}
 
 
-moe_context = {
-    "Qwen3MoeForCausalLM": update_qwen3_moe,
-}
+# Register all MoE models automatically
+for model_class_name, config in MOE_EXPERTS_REPLACEMENTS.items():
+    register_moe_model(model_class_name, config)
 
 
+@contextlib.contextmanager
 def moe_calibration_context(
     model: PreTrainedModel,
-    stack,
-    calibrate_all_experts: bool = False,
+    calibrate_all_experts: bool = True,
 ):
-    # Temporarily updates the MoE modules within the context
-    # Once the context exists, parameter updates persist
+    """
+    Context manager for MoE calibration that temporarily updates MoE modules.
+
+    Args:
+        model: The model to apply MoE calibration to
+        calibrate_all_experts: Whether to calibrate all experts or only routed ones
+
+    Yields:
+        The model with MoE calibration applied
+    """
     cls_name = model.__class__.__name__
-    if cls_name in moe_context:
-        moe_context.get(cls_name)(model, stack, calibrate_all_experts)
+    moe_context = get_moe_context(cls_name)
+
+    if moe_context is None:
+        # No MoE context registered for this model, yield unchanged
+        yield model
+        return
+
+    # Apply MoE calibration
+    moe_context.apply(model, calibrate_all_experts)
+
+    try:
+        yield model
+    finally:
+        # Restore original state
+        moe_context.restore(model)