Remove replace_module_for_calibration

examples/awq/qwen3-vl-30b-a3b-Instruct-example.py

@@ -3,20 +3,16 @@
 from transformers import AutoProcessor, Qwen3VLMoeForConditionalGeneration
 
 from llmcompressor import oneshot
-from llmcompressor.modeling import replace_modules_for_calibration
 from llmcompressor.modifiers.awq import AWQModifier
 from llmcompressor.utils import dispatch_for_generation
 
-# NOTE: Requires a minimum of transformers 4.57.0
-
 MODEL_ID = "Qwen/Qwen3-VL-30B-A3B-Instruct"
 
 # Load model.
 model = Qwen3VLMoeForConditionalGeneration.from_pretrained(
     MODEL_ID, torch_dtype=torch.bfloat16, device_map=None, trust_remote_code=True
 )
 processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
-model = replace_modules_for_calibration(model)
 
 DATASET_ID = "neuralmagic/calibration"
 NUM_CALIBRATION_SAMPLES = 256

examples/quantization_w4a4_fp4/qwen3_vl_moe_w4a4_fp4.py

@@ -3,19 +3,15 @@
 from transformers import AutoProcessor, Qwen3VLMoeForConditionalGeneration
 
 from llmcompressor import oneshot
-from llmcompressor.modeling import replace_modules_for_calibration
 from llmcompressor.modifiers.quantization import QuantizationModifier
 from llmcompressor.utils import dispatch_for_generation
 
-# NOTE: Requires a minimum of transformers 4.57.0
-
 MODEL_ID = "Qwen/Qwen3-VL-235B-A22B-Instruct"
 
 
 # Load model.
 model = Qwen3VLMoeForConditionalGeneration.from_pretrained(MODEL_ID, torch_dtype="auto")
 processor = AutoProcessor.from_pretrained(MODEL_ID)
-model = replace_modules_for_calibration(model)
 
 DATASET_ID = "neuralmagic/calibration"
 NUM_CALIBRATION_SAMPLES = 20

examples/quantization_w8a8_fp8/qwen3_vl_moe_fp8_example.py

@@ -1,7 +1,6 @@
 from transformers import AutoProcessor, Qwen3VLMoeForConditionalGeneration
 
 from llmcompressor import oneshot
-from llmcompressor.modeling import replace_modules_for_calibration
 from llmcompressor.modifiers.quantization import QuantizationModifier
 
 # NOTE: Requires a minimum of transformers 4.57.0
@@ -11,7 +10,6 @@
 # Load model.
 model = Qwen3VLMoeForConditionalGeneration.from_pretrained(MODEL_ID, torch_dtype="auto")
 processor = AutoProcessor.from_pretrained(MODEL_ID)
-model = replace_modules_for_calibration(model)
 
 # Configure the quantization algorithm and scheme.
 # In this case, we:

src/llmcompressor/entrypoints/oneshot.py

@@ -201,22 +201,24 @@ def apply_recipe_modifiers(
         session.reset()
 
         # (Helen INFERENG-661): validate recipe modifiers before initialization
-        session.initialize(
-            model=self.model,
-            start=-1,
-            recipe=self.recipe,
-            recipe_stage=recipe_stage,
-            recipe_args=self.recipe_args.recipe_args,
-            calib_data=calibration_dataloader,
-        )
-        user_pipeline = self.dataset_args.pipeline
-        modifiers = session.lifecycle.recipe.modifiers
-        pipeline = CalibrationPipeline.from_modifiers(modifiers, user=user_pipeline)
         # Apply MoE calibration context for the entire calibration process
         with moe_calibration_context(
             self.model,
             calibrate_all_experts=self.dataset_args.moe_calibrate_all_experts,
         ):
+            session.initialize(
+                model=self.model,
+                start=-1,
+                recipe=self.recipe,
+                recipe_stage=recipe_stage,
+                recipe_args=self.recipe_args.recipe_args,
+                calib_data=calibration_dataloader,
+            )
+            user_pipeline = self.dataset_args.pipeline
+            pipeline = CalibrationPipeline.from_modifiers(
+                session.lifecycle.recipe.modifiers, user=user_pipeline
+            )
+
             pipeline(
                 self.model,
                 calibration_dataloader,

src/llmcompressor/modeling/__init__.py

@@ -18,4 +18,3 @@
 # TODO: add granite4, Qwen3Next
 
 from .fuse import *
-from .prepare import *

src/llmcompressor/modeling/deepseek_v3.py

@@ -68,20 +68,3 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         hidden_states = final_hidden_states.type(hidden_states.dtype).view(*orig_shape)
         hidden_states = hidden_states + self.shared_experts(residuals)
         return hidden_states
-
-
-# Legacy function for backward compatibility
-def replace(
-    config: DeepseekV3Config,
-    module: OriginalDeepseekV3MoE,
-    calibrate_all_experts: bool,
-):
-    """
-    Legacy replacement function.
-    Use CalibrationDeepseekV3MoE instead.
-    """
-    return CalibrationDeepseekV3MoE(
-        module,
-        config,
-        calibrate_all_experts=calibrate_all_experts,
-    )

src/llmcompressor/modeling/llama4.py

@@ -87,16 +87,3 @@ def __init__(self, config: Llama4TextConfig, original: Llama4TextExperts):
             self[i].gate_proj.weight.data = gate_proj.t().contiguous()
             self[i].up_proj.weight.data = up_proj.t().contiguous()
             self[i].down_proj.weight.data = down.t().contiguous()
-
-
-# Legacy function for backward compatibility
-def replace(config: Llama4Config, module: Llama4TextMoe, calibrate_all_experts: bool):
-    """
-    Legacy replacement function.
-    Use SequentialLlama4TextMoe instead.
-    """
-    return SequentialLlama4TextMoe(
-        module,
-        config,
-        calibrate_all_experts=calibrate_all_experts,
-    )

src/llmcompressor/modeling/qwen3_moe.py

@@ -97,20 +97,3 @@ def forward(self, hidden_states: torch.Tensor):
 
     def restore(self, original: torch.nn.Module) -> torch.nn.Module:
         return original
-
-
-# Legacy function for backward compatibility
-def replace(
-    config: Qwen3MoeConfig,
-    module: OriginalQwen3MoeSparseMoeBlock,
-    calibrate_all_experts: bool,
-):
-    """
-    Legacy replacement function.
-    Use CalibrationQwen3MoeSparseMoeBlock instead.
-    """
-    return CalibrationQwen3MoeSparseMoeBlock(
-        module,
-        config,
-        calibrate_all_experts=calibrate_all_experts,
-    )

src/llmcompressor/modeling/qwen3_next_moe.py

@@ -123,13 +123,3 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 
     def restore(self, original: torch.nn.Module) -> torch.nn.Module:
         return original
-
-
-def replace(
-    config,
-    module,
-    calibrate_all_experts,
-):
-    return CalibrationQwen3NextSparseMoeBlock(
-        config=config, original=module, calibrate_all_experts=calibrate_all_experts
-    )

src/llmcompressor/modeling/qwen3_vl_moe.py

@@ -116,15 +116,3 @@ def __init__(self, config, original):
             self[i].gate_proj.weight.data = gate_proj.t().clone().contiguous()
             self[i].up_proj.weight.data = up_proj.t().clone().contiguous()
             self[i].down_proj.weight.data = down.t().clone().contiguous()
-
-
-def replace(
-    config: "Qwen3VLMoeConfig",
-    original: "Qwen3VLMoeTextSparseMoeBlock",
-    calibrate_all_experts: bool,
-):
-    return CalibrateQwen3VLMoeTextSparseMoeBlock(
-        original=original,
-        config=config,
-        calibrate_all_experts=calibrate_all_experts,
-    )