[MODEL] support hrm_text (#2905)

* add missing test_interns1.py

Signed-off-by: ZX-ModelCloud <zx@modelcloud.ai>

* Support multi-tree module layouts and direct meta tensor materialization

Signed-off-by: ZX-ModelCloud <zx@modelcloud.ai>

* support hrm_text

Signed-off-by: ZX-ModelCloud <zx@modelcloud.ai>

* format

Signed-off-by: ZX-ModelCloud <zx@modelcloud.ai>

* cleanup

Signed-off-by: ZX-ModelCloud <zx@modelcloud.ai>

---------

Signed-off-by: ZX-ModelCloud <zx@modelcloud.ai>

Author: ZX-ModelCloud

gptqmodel/__init__.py

@@ -6,6 +6,7 @@
 import os
 import sys
 
+
 # isort: off
 from ._banner import get_startup_banner  # noqa: E402
 from .utils import _MONKEY_PATCH_LOCK  # noqa: E402

gptqmodel/looper/module_looper.py

@@ -51,7 +51,14 @@
     rehome_module_to_device,
     select_forward_devices,
 )
-from ..utils.model import find_modules, get_module, get_module_by_name_prefix, move_to, MoETopKState, set_moe_topk, restore_moe_topk
+from ..utils.model import (
+    MoETopKState,
+    get_layers_with_prefixes,
+    get_module,
+    move_to,
+    restore_moe_topk,
+    set_moe_topk, get_module_by_name_prefix,
+)
 from ..utils.offload import offload_to_disk
 from ..utils.python import has_gil_control, has_gil_disabled
 from ..utils.torch import (CPU, META, timed_gc_collect, torch_sync, tf32_high_precision_guard)
@@ -1428,7 +1435,10 @@ def _loop_impl(self, fallback=None, **kwargs):
 
         forward_pass_use_cache = self.gptq_model.model.config.use_cache if hasattr(self.gptq_model.model.config, "use_cache") else False
         self.gptq_model.model.config.use_cache = False
-        layers, layers_prefix = get_module_by_name_prefix(self.gptq_model.model, self.gptq_model.extract_layers_node())
+        layers, layer_names = get_layers_with_prefixes(
+            self.gptq_model.model,
+            self.gptq_model.extract_layers_node(),
+        )
         region_timer = getattr(self.gptq_model, "quant_region_timer", None)
 
         for p_index, processor in enumerate(self.processors):
@@ -1523,7 +1533,7 @@ def _loop_impl(self, fallback=None, **kwargs):
             layers=layers,
             layer_modules=layer_modules,
             planning_layer_modules=planning_layer_modules,
-            layers_prefix=layers_prefix,
+            layer_names=layer_names,
             fallback=fallback,
             shared_kv_cache_dict=shared_kv_cache_dict,
             pb=pb,
@@ -1631,7 +1641,7 @@ def create_named_modules(self, module, full, is_lm_head_module, layer_index, lay
                 capture_only_flags[n] = True  # forward-only modules should not be finalized
         skipped_modules = []
         for name in subset:
-            layer_name = self.gptq_model.lm_head if is_lm_head_module else f"{layers_prefix}.{layer_index}.{name}"
+            layer_name = self.gptq_model.lm_head if is_lm_head_module else f"{layers_prefix}.{name}"
 
             # gptq task is created and stored inside processor
             if not isinstance(subset[name], NamedModule):

gptqmodel/looper/stage_inputs_capture.py

@@ -19,7 +19,7 @@
 from ..utils.device import get_device
 from ..utils.looper_helpers import device_ctx, select_forward_devices
 from ..utils.logger import setup_logger
-from ..utils.model import get_module_by_name_prefix, move_to, nested_move_to
+from ..utils.model import get_module, get_module_by_name_prefix, move_to, nested_move_to
 from ..utils.torch import CPU, META
 
 if TYPE_CHECKING:  # pragma: no cover - import for typing only
@@ -36,6 +36,15 @@ def __init__(self, looper: ModuleLooper, logger=None) -> None:
         self.gptq_model = looper.gptq_model
         self.logger = logger or setup_logger()
 
+    def _materialize_modules_with_direct_meta_tensors(self, device: torch.device) -> None:
+        for module_name in self.gptq_model.get_modules_with_direct_meta_tensors(self.gptq_model.model):
+            module = get_module(self.gptq_model.model, module_name)
+            if isinstance(module, torch.nn.Module):
+                self.gptq_model.shell_direct_meta_materialize(
+                    target_submodule=module,
+                    device=device,
+                )
+
     def cache_inputs(
         self,
         layers: Sequence[torch.nn.Module],
@@ -174,6 +183,9 @@ def store_input_hook(module, args, kwargs):
             # and wait for the first instance this callback is called
             raise STOP_FORWARD_EXCEPTION
 
+        # Parameters attached to the shell root must be ready before embedding forward.
+        self._materialize_modules_with_direct_meta_tensors(cur_layer_device)
+
         ori_outside_layer_module_devices: Dict[str, torch.device] = {}
         for module_name in self.gptq_model.get_base_modules(self.gptq_model.model):
             module, _ = get_module_by_name_prefix(self.gptq_model.model, [module_name])

gptqmodel/looper/stage_layer.py

@@ -36,7 +36,7 @@
 from ..utils.device import get_device, get_device_new
 from ..utils.looper_helpers import normalize_device_like
 from ..utils.logger import live_renderables_suppressed, log_time_block, setup_logger
-from ..utils.model import find_modules, get_module
+from ..utils.model import find_modules, get_layer_name, get_module
 from ..utils.offload import offload_to_disk
 from ..utils.torch import CPU, torch_empty_cache, torch_sync
 from .stage_subset import SubsetPlan, build_layer_subset_plans, run_subset_stage
@@ -49,11 +49,11 @@ def _find_last_quantized_layer_index(
     looper: "ModuleLooper",
     *,
     layer_modules: List[List[str]],
-    layers_prefix: Optional[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 layers_prefix:
+    if looper.gptq_model.quantize_config.lm_head or not layer_names:
         return None
 
     layer_module_names = {
@@ -67,8 +67,9 @@ def _find_last_quantized_layer_index(
 
     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"{layers_prefix}.{candidate_layer_index}.{module_name}"
+            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:
@@ -387,7 +388,7 @@ def run_layer_stage(
     layers: List[torch.nn.Module],
     layer_modules: List[List[str]],
     planning_layer_modules: List[List[str]],
-    layers_prefix: Optional[str],
+    layer_names: Optional[List[str]],
     fallback,
     shared_kv_cache_dict: Dict[int, torch.Tensor],
     pb,
@@ -403,7 +404,7 @@ def run_layer_stage(
     last_quantized_layer_index = _find_last_quantized_layer_index(
         looper,
         layer_modules=layer_modules,
-        layers_prefix=layers_prefix,
+        layer_names=layer_names,
         layer_count=layer_count,
     )
 
@@ -436,10 +437,12 @@ def run_layer_stage(
             layer_title = "Quantizing lm_head"
             module = get_module(looper.gptq_model.model, key=looper.gptq_model.lm_head)
             pristine_group_module = None
+            layer_name = ""
         else:
             layer_title = f"Quantizing layer {layer_index} of {layer_count - 1}"
             module = layers[layer_index]
             pristine_group_module = None
+            layer_name = get_layer_name(layer_names, layer_index)
 
         pb.title(layer_title).subtitle("").draw()
         if durable_progress_logs:
@@ -483,8 +486,8 @@ def run_layer_stage(
 
             layers[layer_index] = module
 
-            if layers_prefix:
-                layer_descriptor = f"{layers_prefix}.{layer_index}"
+            if layer_name:
+                layer_descriptor = layer_name
             else:
                 layer_descriptor = str(layer_index)
 
@@ -530,7 +533,7 @@ def run_layer_stage(
                 full=full,
                 is_lm_head_module=is_lm_head_module,
                 layer_index=layer_index,
-                layers_prefix=layers_prefix,
+                layers_prefix=layer_name,
                 fallback=fallback,
             )
             if durable_progress_logs:

gptqmodel/looper/weight_only_looper.py

@@ -27,7 +27,14 @@
 from ..nn_modules.converter import MODULE_CONVERTER_MAP
 from ..quantization.config import BitsAndBytesConfig, FP8Config, GGUFConfig, RTNConfig
 from ..utils.logger import setup_logger
-from ..utils.model import find_modules, get_module, get_module_by_name_prefix, move_to
+from ..utils.model import (
+    find_modules,
+    get_layer_name,
+    get_layers_with_prefixes,
+    get_module,
+    get_module_by_name_prefix,
+    move_to,
+)
 from ..utils.offload import offload_to_disk
 
 
@@ -49,7 +56,7 @@ def _resolve_named_module(
         layer_module: torch.nn.Module,
         full: Dict[str, torch.nn.Module],
         layer_index: int,
-        layers_prefix: Optional[str],
+        layer_path: Optional[str],
         module_name: str,
         is_lm_head_module: bool,
     ) -> Optional[NamedModule]:
@@ -65,7 +72,7 @@ def _resolve_named_module(
         if isinstance(resolved, NamedModule):
             return resolved
 
-        layer_name = self.gptq_model.lm_head if is_lm_head_module else f"{layers_prefix}.{layer_index}.{module_name}"
+        layer_name = self.gptq_model.lm_head if is_lm_head_module else f"{layer_path}.{module_name}"
         named = NamedModule(
             resolved,
             name=module_name,
@@ -132,11 +139,19 @@ def loop(self, **kwargs):
         # decoder-cache state while layers are being replaced.
         self.gptq_model.model.config.use_cache = False
 
-        layers, layers_prefix = get_module_by_name_prefix(
+        layers, layer_names = get_layers_with_prefixes(
             self.gptq_model.model,
             self.gptq_model.extract_layers_node(),
         )
 
+        for module_name in self.gptq_model.get_modules_with_direct_meta_tensors(self.gptq_model.model):
+            module = get_module(self.gptq_model.model, module_name)
+            if module is not None:
+                self.gptq_model.shell_direct_meta_materialize(
+                    target_submodule=module,
+                    device=CPU,
+                )
+
         if quant_config.offload_to_disk:
             log.info("Offloading base modules to disk...")
             offload_to_disk(
@@ -181,6 +196,10 @@ def loop(self, **kwargs):
                 else:
                     module = layers[layer_index]
                     subsets = layer_modules
+                    # Flattened layer names preserve the source stack for split decoders.
+                    layer_name = get_layer_name(layer_names, layer_index)
+                if is_lm_head_module:
+                    layer_name = None
 
                 module = self.gptq_model.pre_quantize(module)
                 if not is_lm_head_module:
@@ -204,7 +223,7 @@ def loop(self, **kwargs):
                             layer_module=module,
                             full=full,
                             layer_index=layer_index,
-                            layers_prefix=layers_prefix,
+                            layer_path=layer_name,
                             module_name=module_name,
                             is_lm_head_module=is_lm_head_module,
                         )

gptqmodel/models/auto.py

@@ -113,6 +113,7 @@
 from .definitions.gptj import GptJQModel  # noqa: E402
 from .definitions.granitemoehybrid import GraniteMoeHybridQModel
 from .definitions.grinmoe import GrinMoeQModel  # noqa: E402
+from .definitions.hrm_text import HrmTextQModel  # noqa: E402
 from .definitions.hymba import HymbaQModel  # noqa: E402
 from .definitions.instella import InstellaQModel  # noqa: E402
 from .definitions.internlm import InternLMQModel  # noqa: E402
@@ -227,6 +228,7 @@
     "internlm2": InternLM2QModel,
     "interns1": InternS1QModel,
     "internvl_chat": InternVLChatQModel,
+    "hrm_text": HrmTextQModel,
     "qwen": QwenQModel,
     "mistral": LlamaQModel, # 100% llama clone
     "yi": LlamaQModel, # 100% llama clone