Merge branch 'main' into fix/granite4-example-updates

Author: brian-dellabetta

src/llmcompressor/modeling/qwen3_next_moe.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 # coding=utf-8
 # Copyright 2025 The Qwen team, Alibaba Group and the HuggingFace Inc. team.
 # All rights reserved.
@@ -13,19 +15,21 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from typing import TYPE_CHECKING
 
 import torch
 
 from llmcompressor.modeling.moe_context import MoECalibrationModule
 
-
-@MoECalibrationModule.register("Qwen3NextSparseMoeBlock")
-class CalibrationQwen3NextSparseMoeBlock(MoECalibrationModule):
+if TYPE_CHECKING:
     from transformers import Qwen3NextConfig
     from transformers.models.qwen3_next.modeling_qwen3_next import (
         Qwen3NextSparseMoeBlock,
     )
 
+
+@MoECalibrationModule.register("Qwen3NextSparseMoeBlock")
+class CalibrationQwen3NextSparseMoeBlock(MoECalibrationModule):
     """
     Calibration version of Qwen3NextSparseMoeBlock that sends all tokens to all experts.
     """

tests/e2e/e2e_utils.py

@@ -10,15 +10,12 @@
 from tests.testing_utils import process_dataset
 
 
-@log_time
-def _load_model_and_processor(
-    model: str,
-    model_class: str,
-):
+def load_model(model: str, model_class: str, device_map: str | None = None):
     pretrained_model_class = getattr(transformers, model_class)
-    loaded_model = pretrained_model_class.from_pretrained(model, torch_dtype="auto")
-    processor = AutoProcessor.from_pretrained(model)
-    return loaded_model, processor
+    loaded_model = pretrained_model_class.from_pretrained(
+        model, torch_dtype="auto", device_map=device_map
+    )
+    return loaded_model
 
 
 @log_time
@@ -41,9 +38,8 @@ def run_oneshot_for_e2e_testing(
     # Load model.
     oneshot_kwargs = {}
 
-    loaded_model, processor = _load_model_and_processor(
-        model=model, model_class=model_class
-    )
+    loaded_model = load_model(model=model, model_class=model_class)
+    processor = AutoProcessor.from_pretrained(model)
 
     if dataset_id:
         ds = load_dataset(dataset_id, name=dataset_config, split=dataset_split)
@@ -88,11 +84,13 @@ def data_collator(batch):
                 targets="Linear",
                 scheme=scheme,
                 actorder=None,  # added for consistency with past testing configs
-                ignore=["lm_head"],
+                ignore=["lm_head", "re:.*mlp.gate[.].*"],
             )
         else:
             oneshot_kwargs["recipe"] = QuantizationModifier(
-                targets="Linear", scheme=scheme, ignore=["lm_head"]
+                targets="Linear",
+                scheme=scheme,
+                ignore=["lm_head", "re:.*mlp.gate[.].*"],
             )
 
     # Apply quantization.

tests/e2e/vLLM/configs/qwen3_fp4_nvfp4.yaml

@@ -0,0 +1,7 @@
+cadence: "nightly"
+test_type: "regression"
+model: Qwen/Qwen3-30B-A3B
+scheme: NVFP4
+dataset_id: HuggingFaceH4/ultrachat_200k
+dataset_split: train_sft
+num_calibration_samples: 20

tests/e2e/vLLM/configs/qwen3_fp8_dynamic_per_token.yaml

@@ -0,0 +1,4 @@
+cadence: "nightly"
+test_type: "regression"
+model: Qwen/Qwen3-30B-A3B
+scheme: FP8_DYNAMIC

tests/e2e/vLLM/run_vllm.py

@@ -18,7 +18,7 @@ def parse_args():
     except json.JSONDecodeError as e:
         raise ValueError(f"Invalid JSON input: {e}")
 
-    if "W4A16_2of4" in scheme:
+    if scheme is not None and "W4A16_2of4" in scheme:
         # required by the kernel
         llm_kwargs["dtype"] = torch.float16
 

tests/llmcompressor/modeling/test_calib_qwen3_next.py

@@ -11,9 +11,22 @@
 from llmcompressor.utils.helpers import DisableQuantization, calibration_forward_context
 from tests.testing_utils import requires_cadence, requires_gpu
 
+try:
+    from transformers import Qwen3NextConfig
+    from transformers.models.qwen3_next.modeling_qwen3_next import (
+        Qwen3NextSparseMoeBlock,
+    )
+except ImportError:
+    Qwen3NextConfig = None
+    Qwen3NextSparseMoeBlock = None
+
 
 @requires_cadence("weekly")
 @pytest.mark.parametrize("model_stub", ["Qwen/Qwen3-Next-80B-A3B-Instruct"])
+@pytest.mark.skipif(
+    Qwen3NextConfig is None,
+    reason="Qwen3Next not available in this version of transformers",
+)
 def test_calib_replace_qwen3moe_all_experts(model_stub):
     with skip_weights_download():
         model = AutoModelForCausalLM.from_pretrained(model_stub)
@@ -60,12 +73,11 @@ def hook_fn(i, module, input, output):
 
 
 @requires_gpu
+@pytest.mark.skipif(
+    Qwen3NextConfig is None,
+    reason="Qwen3Next not available in this version of transformers",
+)
 def test_calib_qwen3_moe_module():
-    from transformers import Qwen3NextConfig
-    from transformers.models.qwen3_next.modeling_qwen3_next import (
-        Qwen3NextSparseMoeBlock,
-    )
-
     config = Qwen3NextConfig()
     with torch.device("cuda"):
         original = Qwen3NextSparseMoeBlock(config).eval()

tests/lmeval/test_lmeval.py

@@ -13,7 +13,7 @@
 from pydantic import BaseModel
 
 from llmcompressor.core import active_session
-from tests.e2e.e2e_utils import run_oneshot_for_e2e_testing
+from tests.e2e.e2e_utils import load_model, run_oneshot_for_e2e_testing
 from tests.test_timer.timer_utils import get_singleton_manager, log_time
 from tests.testing_utils import requires_gpu
 
@@ -35,6 +35,10 @@ class LmEvalConfig(BaseModel):
 
 try:
     import lm_eval
+    import lm_eval.api.registry
+
+    # needed to populate model registry
+    import lm_eval.models  # noqa
 
     lm_eval_installed = True
 except ImportError:
@@ -120,7 +124,7 @@ def test_lm_eval(self, test_data_file: str):
 
         # Always evaluate base model for recovery testing
         logger.info("================= Evaluating BASE model ======================")
-        self.base_results = self._eval_base_model()
+        base_results = self._eval_base_model()
 
         if not self.save_dir:
             self.save_dir = self.model.split("/")[1] + f"-{self.scheme}"
@@ -145,22 +149,41 @@ def test_lm_eval(self, test_data_file: str):
         self._handle_recipe()
 
         logger.info("================= Running LM Eval on COMPRESSED model ==========")
-        self._run_lm_eval()
+        compressed_results = self._eval_compressed_model()
+
+        # Always use recovery testing
+        self._validate_recovery(base_results, compressed_results)
+
+        # If absolute metrics provided, show warnings (not failures)
+        if self.lmeval.metrics:
+            self._check_absolute_warnings(compressed_results)
 
         self.tear_down()
 
     @log_time
-    def _eval_base_model(self):
+    def _eval_base_model(self) -> dict:
         """Evaluate the base (uncompressed) model."""
-        model_args = {**self.lmeval.model_args, "pretrained": self.model}
+        return self._eval_model(self.model)
+
+    @log_time
+    def _eval_compressed_model(self) -> dict:
+        """Evaluate the compressed model."""
+        return self._eval_model(self.save_dir)
+
+    def _eval_model(self, model: str) -> dict:
+        # NOTE: pass in PreTrainedModel to avoid lm_eval's model-loading logic
+        # https://github.com/EleutherAI/lm-evaluation-harness/pull/3393
+        lm_eval_cls = lm_eval.api.registry.get_model(self.lmeval.model)
 
         results = lm_eval.simple_evaluate(
-            model=self.lmeval.model,
-            model_args=model_args,
+            model=lm_eval_cls(
+                pretrained=load_model(model, self.model_class, device_map="cuda:0"),
+                batch_size=self.lmeval.batch_size,
+                **self.lmeval.model_args,
+            ),
             tasks=[self.lmeval.task],
             num_fewshot=self.lmeval.num_fewshot,
             limit=self.lmeval.limit,
-            device="cuda:0",
             apply_chat_template=self.lmeval.apply_chat_template,
             batch_size=self.lmeval.batch_size,
         )
@@ -181,31 +204,9 @@ def _handle_recipe(self):
             fp.write(recipe_yaml_str)
         session.reset()
 
-    @log_time
-    def _run_lm_eval(self):
-        model_args = {"pretrained": self.save_dir}
-        model_args.update(self.lmeval.model_args)
-        results = lm_eval.simple_evaluate(
-            model=self.lmeval.model,
-            model_args=model_args,
-            tasks=[self.lmeval.task],
-            num_fewshot=self.lmeval.num_fewshot,
-            limit=self.lmeval.limit,
-            device="cuda:0",
-            apply_chat_template=self.lmeval.apply_chat_template,
-            batch_size=self.lmeval.batch_size,
-        )
-
-        # Always use recovery testing
-        self._validate_recovery(results)
-
-        # If absolute metrics provided, show warnings (not failures)
-        if self.lmeval.metrics:
-            self._check_absolute_warnings(results)
-
-    def _validate_recovery(self, compressed_results):
+    def _validate_recovery(self, base_results, compressed_results):
         """Validate using recovery testing - compare against base model."""
-        base_metrics = self.base_results["results"][self.lmeval.task]
+        base_metrics = base_results["results"][self.lmeval.task]
         compressed_metrics = compressed_results["results"][self.lmeval.task]
         higher_is_better_map = compressed_results.get("higher_is_better", {}).get(
             self.lmeval.task, {}