Merge branch 'main' into kylesayrs/deprecate

Author: dsikka

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:

setup.py

@@ -114,22 +114,28 @@ def localversion_func(version: ScmVersion) -> str:
         ("pyyaml>=6.0.1,<=6.0.3" if BUILD_TYPE == "release" else "pyyaml>=6.0.1"),
         # librosa dependency numba is currently not compatible with numpy>=2.3
         # https://numba.readthedocs.io/en/stable/user/installing.html#version-support-information
-        ("numpy>=2.0.0,<=2.3.3" if BUILD_TYPE == "release" else "numpy>=2.0.0"),
+        ("numpy>=2.0.0,<=2.3.5" if BUILD_TYPE == "release" else "numpy>=2.0.0"),
         (
             "requests>=2.32.2,<=2.32.5"
             if BUILD_TYPE == "release"
             else "requests>=2.32.2"
         ),
         ("tqdm>=4.66.3,<=4.67.1" if BUILD_TYPE == "release" else "tqdm>=4.66.3"),
-        ("torch>=2.7.0,<=2.8.0" if BUILD_TYPE == "release" else "torch>=2.7.0"),
+        ("torch>=2.7.0,<=2.9.1" if BUILD_TYPE == "release" else "torch>=2.7.0"),
         (
-            "transformers>=4.54.0,<=4.56.2"
+            "transformers>=4.54.0,<=4.57.3"
             if BUILD_TYPE == "release"
-            else "transformers>=4.54.0"
+            else "transformers>=4.54.0,<=4.57.3"
         ),
-        ("datasets>=4.0.0,<=4.1.1" if BUILD_TYPE == "release" else "datasets>=4.0.0"),
+        ("datasets>=4.0.0,<=4.4.1" if BUILD_TYPE == "release" else "datasets>=4.0.0"),
         (
-            "accelerate>=1.6.0,<=1.10.1"
+            # auto-round 0.9.1 cannot work with accelerate <1.10.0
+            "auto-round>=0.9.2,<=0.9.2"
+            if BUILD_TYPE == "release"
+            else "auto-round>=0.9.2"
+        ),
+        (
+            "accelerate>=1.6.0,<=1.12.0"
             if BUILD_TYPE == "release"
             else "accelerate>=1.6.0"
         ),
@@ -138,13 +144,12 @@ def localversion_func(version: ScmVersion) -> str:
             if BUILD_TYPE == "release"
             else "nvidia-ml-py>=12.560.30"
         ),
-        ("pillow>=10.4.0,<=11.3.0" if BUILD_TYPE == "release" else "pillow>=10.4.0"),
+        ("pillow>=10.4.0,<=12.0.0" if BUILD_TYPE == "release" else "pillow>=10.4.0"),
         (
             "compressed-tensors==0.12.2"
             if BUILD_TYPE == "release"
             else "compressed-tensors>=0.12.3a2"
         ),
-        ("auto-round==0.9.1"),
     ],
     extras_require={
         "dev": [
@@ -178,7 +183,7 @@ def localversion_func(version: ScmVersion) -> str:
             "mkdocstrings-python",
             "mkdocs-gen-files",
             "mkdocs-nav-weight",
-        ]
+        ],
     },
     entry_points={
         "console_scripts": [

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/prepare.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,
-    )