Fix the issue where WeightOnlyLooper did not support dynamic skipping.

gptqmodel/looper/stage_layer.py

@@ -34,7 +34,7 @@
 from ..looper.paroquant_processor import ParoQuantProcessor
 from ..looper.qqq_processor import QQQProcessor
 from ..utils.device import get_device, get_device_new
-from ..utils.looper_helpers import normalize_device_like
+from ..utils.looper_helpers import find_last_quantized_layer_index, normalize_device_like
 from ..utils.logger import live_renderables_suppressed, log_time_block, setup_logger
 from ..utils.model import find_modules, get_layer_name, get_module
 from ..utils.offload import offload_to_disk
@@ -45,40 +45,6 @@
     from .module_looper import ModuleLooper
 
 
-def _find_last_quantized_layer_index(
-    looper: "ModuleLooper",
-    *,
-    layer_modules: List[List[str]],
-    layer_names: Optional[List[str]],
-    layer_count: int,
-) -> Optional[int]:
-    """Return the highest layer index whose tracked modules are not all dynamically skipped."""
-    if looper.gptq_model.quantize_config.lm_head or not layer_names:
-        return None
-
-    layer_module_names = {
-        name.split("#", 1)[0]
-        for module_group in layer_modules
-        for name in module_group
-        if name
-    }
-    if not layer_module_names:
-        return None
-
-    last_quantized_layer_index = -1
-    for candidate_layer_index in range(layer_count):
-        layer_name = get_layer_name(layer_names, candidate_layer_index)
-        for module_name in layer_module_names:
-            module_full_name = f"{layer_name}.{module_name}"
-            # If at least one module in this layer is not dynamically excluded,
-            # the layer still needs forward/quantization work.
-            if looper.gptq_model.quantize_config.dynamic_get(layer_name=module_full_name) != False:
-                last_quantized_layer_index = candidate_layer_index
-                break
-
-    return last_quantized_layer_index
-
-
 def _should_drain_finalize_futures_synchronously(
     looper: "ModuleLooper",
     *,
@@ -401,8 +367,8 @@ def run_layer_stage(
     # Trailing layers whose tracked modules are all dynamically excluded never
     # need another forward or finalize pass, so the loop can stop once the
     # final eligible layer has been processed.
-    last_quantized_layer_index = _find_last_quantized_layer_index(
-        looper,
+    last_quantized_layer_index = find_last_quantized_layer_index(
+        looper.gptq_model.quantize_config,
         layer_modules=layer_modules,
         layer_names=layer_names,
         layer_count=layer_count,

gptqmodel/looper/weight_only_looper.py

@@ -35,7 +35,13 @@
 from ..utils.device import get_device
 from ..utils.device_telemetry import emit_device_telemetry
 from ..utils.logger import log_time_block, setup_logger
-from ..utils.looper_helpers import device_ctx, normalize_device_like, rehome_module_to_device, select_forward_devices
+from ..utils.looper_helpers import (
+    device_ctx,
+    find_last_quantized_layer_index,
+    normalize_device_like,
+    rehome_module_to_device,
+    select_forward_devices,
+)
 from ..utils.model import (
     find_modules,
     get_layer_name,
@@ -756,9 +762,34 @@ def loop(self, **kwargs):
             )
 
         try:
+            # Trailing layers whose tracked modules are all dynamically excluded never
+            # need another forward or finalize pass, so the loop can stop once the
+            # final eligible layer has been processed.
+            last_quantized_layer_index = find_last_quantized_layer_index(
+                self.gptq_model.quantize_config,
+                layer_modules=layer_modules,
+                layer_names=layer_names,
+                layer_count=layer_count,
+            )
+
             for layer_index in range(total_layers):
                 is_lm_head_module = layer_index >= layer_count
 
+                if (
+                    not is_lm_head_module
+                    and last_quantized_layer_index is not None
+                    and layer_index > last_quantized_layer_index
+                ):
+                    # The remaining layers are fully skipped by dynamic config, so
+                    # avoid entering another layer-level quantization cycle.
+                    log.debug(
+                        "StageLayer: early stop at layer=%s, last_quantized_layer=%s",
+                        layer_index,
+                        last_quantized_layer_index,
+                    )
+                    pb.close()
+                    break
+
                 # Transformer blocks and lm_head follow the same weight-only
                 # lifecycle, but lm_head is resolved from the root model.
                 if is_lm_head_module:
@@ -816,6 +847,12 @@ def loop(self, **kwargs):
                         if named is None:
                             continue
 
+                        # Match ModuleLooper's processor lifecycle: dynamically
+                        # excluded modules must not enter device scheduling or
+                        # quantization work.
+                        if self.processor.is_skipped(named):
+                            continue
+
                         preferred_device = layer_strategy_device_map.get(module_name)
                         if preferred_device is not None:
                             # Weight-only has no SubsetPlan, so store the same

gptqmodel/looper/weight_only_processor.py

@@ -68,6 +68,11 @@ def __init__(
         )
         self.lock = threading.Lock()
 
+    def is_skipped(self, module: NamedModule) -> bool:
+        """Report whether dynamic configuration excludes this module."""
+
+        return self.qcfg.dynamic_get(layer_name=module.full_name) is False
+
     @staticmethod
     def _uses_direct_pack(qcfg: RTNConfig | GGUFConfig | FP8Config | BitsAndBytesConfig) -> bool:
         """Returns whether the method packs directly from the original dense weights."""

gptqmodel/utils/looper_helpers.py

@@ -20,7 +20,7 @@
 from ..utils.env import env_flag
 from ..utils.inspect import get_supported_kwargs
 from ..utils.logger import setup_logger
-from ..utils.model import move_to, nested_move_to
+from ..utils.model import get_layer_name, move_to, nested_move_to
 from ..utils.safe import ThreadSafe
 from ..utils.torch import ALL_DEVICES, CPU, HAS_NPU, torch_sync
 
@@ -65,6 +65,7 @@ def torch_replicate(
 if TYPE_CHECKING:
     from ..looper.loop_processor import LoopProcessor
     from ..models._const import DEVICE
+    from ..quantization.config import BaseQuantizeConfig
 
 
 __all__ = [
@@ -74,10 +75,44 @@ def torch_replicate(
     "select_forward_devices",
     "normalize_device_like",
     "clone_module_for_devices",
+    "find_last_quantized_layer_index",
     "forward_batch_worker",
 ]
 
 
+def find_last_quantized_layer_index(
+    quantize_config: "BaseQuantizeConfig",
+    *,
+    layer_modules: List[List[str]],
+    layer_names: Optional[List[str]],
+    layer_count: int,
+) -> Optional[int]:
+    """Return the final layer containing a module not excluded by dynamic config."""
+
+    if quantize_config.lm_head or not layer_names:
+        return None
+
+    layer_module_names = {
+        name.split("#", 1)[0]
+        for module_group in layer_modules
+        for name in module_group
+        if name
+    }
+    if not layer_module_names:
+        return None
+
+    last_quantized_layer_index = -1
+    for candidate_layer_index in range(layer_count):
+        layer_name = get_layer_name(layer_names, candidate_layer_index)
+        if any(
+            quantize_config.dynamic_get(layer_name=f"{layer_name}.{module_name}") is not False
+            for module_name in layer_module_names
+        ):
+            last_quantized_layer_index = candidate_layer_index
+
+    return last_quantized_layer_index
+
+
 @contextmanager
 def device_ctx(dev: Optional[torch.device | "DEVICE"]):
     """Temporarily set the thread-local device for CUDA/XPU backends."""

tests/test_looper_helpers.py

@@ -12,6 +12,22 @@ def _set_current_batch_index(self, batch_index):
         self.current_batch_index = batch_index
 
 
+class _DynamicConfig:
+    lm_head = False
+
+    def dynamic_get(self, *, layer_name):
+        return False if layer_name.startswith(("layers.1.", "layers.2.")) else None
+
+
+def test_find_last_quantized_layer_index_uses_dynamic_exclusions():
+    assert looper_helpers.find_last_quantized_layer_index(
+        _DynamicConfig(),
+        layer_modules=[["linear#capture_only"]],
+        layer_names=["layers.0", "layers.1", "layers.2"],
+        layer_count=3,
+    ) == 0
+
+
 class _RequiresAttentionMask(torch.nn.Module):
     def __init__(self):
         super().__init__()

tests/test_stage_modules.py

@@ -973,6 +973,9 @@ def __init__(self):
         def pre_quantize(self, module):
             return module
 
+        def should_quantize_layer(self, *_args):
+            return True
+
         def post_quantize(self, module):
             return module
 
@@ -1043,7 +1046,7 @@ def create_named_modules(self, module, full, is_lm_head_module, layer_index, lay
         layers=[torch.nn.Linear(64, 64) for _ in range(3)],
         layer_modules=[["foo"]],
         planning_layer_modules=[["foo"]],
-        layers_prefix="model.layers",
+        layer_names=["model.layers.0", "model.layers.1", "model.layers.2"],
         fallback=True,
         shared_kv_cache_dict={},
         pb=pb,