From 3b0a054ac1ba03c8ba50744160da399128e6956d Mon Sep 17 00:00:00 2001 From: ZX-ModelCloud Date: Mon, 31 Aug 2026 10:21:30 +0800 Subject: [PATCH] feat: support Qwen3.8 Flash Next quantization --- README.md | 3 +- gptqmodel/looper/stage_subset.py | 51 ++- gptqmodel/models/auto.py | 2 + gptqmodel/models/base.py | 41 +++ gptqmodel/models/definitions/__init__.py | 1 + gptqmodel/models/definitions/qwen4_exp.py | 75 ++++ gptqmodel/models/loader.py | 17 +- gptqmodel/utils/model.py | 106 +++++- gptqmodel/utils/structure.py | 84 ++++- tests/models/test_qwen4_exp.py | 39 +++ tests/test_lazy_turtle_conversion_mapping.py | 75 +++- tests/test_qwen4_exp_support.py | 338 +++++++++++++++++++ tests/test_subset_plan.py | 65 ++++ 13 files changed, 857 insertions(+), 40 deletions(-) create mode 100644 gptqmodel/models/definitions/qwen4_exp.py create mode 100644 tests/models/test_qwen4_exp.py create mode 100644 tests/test_qwen4_exp_support.py diff --git a/README.md b/README.md index 3397f0fc6..d8f034b62 100644 --- a/README.md +++ b/README.md @@ -21,6 +21,7 @@ ## Latest News 🗞️🚀 +* 08/31/2026 7.4.0-dev `main`: ✨ Added Qwen3.8-Flash-Next (`qwen4_exp`) quantization. * 08/26/2026 7.4.0-dev `main`: ✨ Added NVIDIA `LocateAnything-3B` quantization support. * 08/25/2026 7.4.0-dev `main`: ✨ Added Tencent `HunyuanOCR` quantization support. * 08/25/2026 7.4.0-dev `main`: ✨ Added `lm_head` and embedding quantization lifecycle. @@ -255,7 +256,7 @@ Selected public references where teams or companies explicitly mention GPT-QMode | Model | | | | | | | | | | |-------------------------------|---|---------------------------------|--|----------------------------|--|---------------------------------|--|------------------------|---| -| Apertus | ✅ | EXAONE 3/4 | ✅ | Dots1 | ✅ | Mistral3 / Ministral3 | ✅ | Qwen 2/3/3.5 (Next/MoE) | ✅ | +| Apertus | ✅ | EXAONE 3/4 | ✅ | Dots1 | ✅ | Mistral3 / Ministral3 | ✅ | Qwen 2/3/3.5/3.8 (Next/MoE) | ✅ | | Baichuan | ✅ | Falcon (H1 / Mamba) | ✅ | InternLM 1/2/2.5 | ✅ | Mixtral | ✅ | Qwen 2/2.5/3 VL | ✅ | | Bloom | ✅ | FastVLM | ✅ | Kimi K2 | ✅ | MobileLLM | ✅ | Qwen 2.5/3 Omni | ✅ | | ChatGLM | ✅ | Gemma 1-4 / 3n | ✅ | Klear | ✅ | MOSS | ✅ | RefinedWeb | ✅ | diff --git a/gptqmodel/looper/stage_subset.py b/gptqmodel/looper/stage_subset.py index 0a9350cdc..55948033e 100644 --- a/gptqmodel/looper/stage_subset.py +++ b/gptqmodel/looper/stage_subset.py @@ -498,24 +498,41 @@ def build_subset_plan( for module_name in moe_groups[group_key]: forward_device_map[module_name] = target_device - if forward_device_map: - # Once either dense or expert placement is explicit, anchor every - # untouched module back to its baseline placement so stale quant - # devices never leak into a later subset forward. - baseline_devices = _resolve_forward_baseline_devices( - subset=subset, - full=full, - ) - for module_name, baseline_device in baseline_devices.items(): - forward_device_map.setdefault(module_name, baseline_device) - - for module_name, named_module in subset.items(): - preferred_device = forward_device_map.get(module_name) - if preferred_device is not None: - named_module.state["preferred_quant_device"] = preferred_device + # A model may keep selected leaf modules on CPU even while replaying their layer on GPU. + placement_override = getattr(looper.gptq_model, "forward_device_for_module", None) + placement_override_active = getattr(looper.gptq_model, "has_forward_device_overrides", None) + placement_override_active = ( + callable(placement_override) + and callable(placement_override_active) + and placement_override_active() + ) - restore_forward_device_overrides = False - subset_forward_serial = True + if forward_device_map or placement_override_active: + # Start from each leaf's current device so an excluded tensor is never + # moved implicitly with its parent layer. + baseline_devices = _resolve_forward_baseline_devices( + subset=subset, + full=full, + ) + for module_name, baseline_device in baseline_devices.items(): + forward_device_map.setdefault(module_name, baseline_device) + + if placement_override_active: + for module_name, planned_device in list(forward_device_map.items()): + module_ref = subset.get(module_name) + if module_ref is None and full is not None: + module_ref = full.get(module_name) + actual_module = module_ref.module if isinstance(module_ref, NamedModule) else module_ref + if actual_module is not None: + forward_device_map[module_name] = placement_override(actual_module, planned_device) + + for module_name, named_module in subset.items(): + preferred_device = forward_device_map.get(module_name) + if preferred_device is not None: + named_module.state["preferred_quant_device"] = preferred_device + + restore_forward_device_overrides = False + subset_forward_serial = True auto_forward_data_parallel = getattr( looper.gptq_model.quantize_config, diff --git a/gptqmodel/models/auto.py b/gptqmodel/models/auto.py index 54a2dfa8c..c1ac7de67 100644 --- a/gptqmodel/models/auto.py +++ b/gptqmodel/models/auto.py @@ -184,6 +184,7 @@ from .definitions.qwen3_next import Qwen3NextGPTQ # noqa: E402 from .definitions.qwen3_omni_moe import Qwen3OmniMoeGPTQ from .definitions.qwen3_vl import Qwen3_VLQModel +from .definitions.qwen4_exp import Qwen4ExpQModel # noqa: E402 from .definitions.rw import RwgQModel # noqa: E402 from .definitions.solar_open import SolarOpenQModel # noqa: E402 from .definitions.solar_open2 import SolarOpen2QModel # noqa: E402 @@ -309,6 +310,7 @@ "qwen2_5_omni": Qwen2_5_OmniGPTQ, "qwen3_omni_moe": Qwen3OmniMoeGPTQ, "qwen3_vl": Qwen3_VLQModel, + "qwen4_exp": Qwen4ExpQModel, "dbrx": DbrxQModel, "dbrx_converted": DbrxConvertedQModel, "deepseek_v2": DeepSeekV2QModel, diff --git a/gptqmodel/models/base.py b/gptqmodel/models/base.py index 913cd1ca6..67ab22b05 100644 --- a/gptqmodel/models/base.py +++ b/gptqmodel/models/base.py @@ -291,6 +291,9 @@ class BaseQModel(nn.Module): # so `defuser_module_paths` is used to explicitly locate and defuse them. defuser_module_paths = None + # Multimodal wrappers can reuse the checkpoint rules of their text model. + hf_conversion_model_type_alias: Optional[str] = None + def __init__( self, model: PreTrainedModel, @@ -488,6 +491,15 @@ def resolve_hf_conversion_map_reversed(cls, target_model: Optional[nn.Module] = if configured_map is not None: return copy.deepcopy(configured_map) + model_type_alias = getattr(cls, "hf_conversion_model_type_alias", None) + if model_type_alias: + inferred_map = LazyTurtle.infer_hf_conversion_map_reversed( + target_model=target_model, + model_type=model_type_alias, + ) + if inferred_map is not None: + return copy.deepcopy(inferred_map) + inferred_map = LazyTurtle.infer_hf_conversion_map_reversed(target_model=target_model) return copy.deepcopy(inferred_map) if inferred_map is not None else None @@ -1973,6 +1985,35 @@ def pre_quantize(self, module: nn.Module) -> nn.Module: else: return module + def forward_device_for_module(self, module: nn.Module, planned_device: torch.device) -> torch.device: + """Apply model-declared placement exclusions to subset replay planning.""" + + turtle_model = self.turtle_model + if not isinstance(turtle_model, LazyTurtle): + return planned_device + + # LazyTurtle matches exclusions by dotted parameter path, not module type. + module_paths = getattr(self, "_forward_module_paths_by_id", None) + if module_paths is None or id(module) not in module_paths: + module_paths = {id(candidate): name for name, candidate in self.model.named_modules() if name} + self._forward_module_paths_by_id = module_paths + module_path = module_paths.get(id(module)) + if module_path is None: + return planned_device + # Check only tensors owned by this leaf; descendants receive their own plan entry. + for rel_name, _ in module.named_parameters(recurse=False): + if turtle_model.is_no_placement_tensor(module_path, rel_name): + return torch.device(CPU) + return planned_device + + def has_forward_device_overrides(self) -> bool: + """Return whether replay must preserve model-declared tensor placement.""" + + turtle_model = self.turtle_model + return isinstance(turtle_model, LazyTurtle) and bool( + getattr(turtle_model, "_no_placement_params", ()) + ) + def post_quantize(self, module: nn.Module) -> nn.Module: #return self.offload_to_disk(module=module) return move_to(module, device=CPU) diff --git a/gptqmodel/models/definitions/__init__.py b/gptqmodel/models/definitions/__init__.py index 41f65b93c..1df76a3f8 100644 --- a/gptqmodel/models/definitions/__init__.py +++ b/gptqmodel/models/definitions/__init__.py @@ -96,6 +96,7 @@ from .qwen3 import Qwen3QModel from .qwen3_moe import Qwen3MoeQModel from .qwen3_vl import Qwen3_VLQModel +from .qwen4_exp import Qwen4ExpQModel from .rw import RwgQModel from .solar_open import SolarOpenQModel from .solar_open2 import SolarOpen2QModel diff --git a/gptqmodel/models/definitions/qwen4_exp.py b/gptqmodel/models/definitions/qwen4_exp.py new file mode 100644 index 000000000..8ecad58b3 --- /dev/null +++ b/gptqmodel/models/definitions/qwen4_exp.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: 2026 ModelCloud.ai +# SPDX-FileCopyrightText: 2026 qubitium@modelcloud.ai +# SPDX-License-Identifier: Apache-2.0 +# Contact: qubitium@modelcloud.ai, x.com/qubitium + +from transformers import AutoModelForImageTextToText + +from ..base import BaseQModel +from ..moe_lifecycle import GateUpDownMoELifecycleHooks + + +class Qwen4ExpQModel(BaseQModel): + """Qwen3.8-Flash-Next / Qwen4 experimental multimodal MoE.""" + + loader = AutoModelForImageTextToText + require_load_processor = True + layer_modules_strict = False + + dynamic_expert_index = "num_experts" + + # The final mixer replaces the usual decoder RMSNorm. + pre_lm_head_norm_module = "model.language_model.hyper_connection_mixer" + rotary_embedding = "model.language_model.rotary_emb" + + # Transformers intentionally ignores the auxiliary MTP decoder on load. + out_of_model_tensors = {"prefixes": ["mtp"]} + + # The outer config reuses the text model's PLE checkpoint mapping. + hf_conversion_model_type_alias = "qwen4_exp_text" + moe_lifecycle_hooks = GateUpDownMoELifecycleHooks() + + # GQA makes o_proj shape-incompatible with the Q/K/V AWQ scale group. + awq_scale_optimize_shape_dependent_modules = ["self_attn.o_proj"] + + # Only unmarked entries are quantized; PLE and hyper weights are omitted. + module_tree = [ + "model", + "language_model", + "layers", + "#", + { + "self_attn": ( + "indexer.index_qk_proj:!", + "indexer.q_layernorm:!", + "indexer.k_layernorm:!", + "q_proj:0", + "q_norm:!", + "k_proj:0", + "k_norm:!", + "v_proj:0", + "o_proj:1", + ), + "linear_attn": ( + "conv1d:!", + "in_proj_qkv:0", + "in_proj_z:1", + "in_proj_b:!:1", + "in_proj_a:!:1", + "norm:!", + "out_proj:2", + ), + "mlp:moe": { + # Keep shared experts separate so placeholder expansion does not duplicate them. + "shared_expert": ("gate_proj:0", "up_proj:0", "down_proj:1"), + "gate": ("gate:!",), + "experts:0": { + "#": ("gate_proj:0", "up_proj:0", "down_proj:1"), + }, + "shared_expert_gate": ("shared_expert_gate:!",), + }, + }, + ] + + +__all__ = ["Qwen4ExpQModel"] diff --git a/gptqmodel/models/loader.py b/gptqmodel/models/loader.py index 70eb8a7f0..2a17f5219 100644 --- a/gptqmodel/models/loader.py +++ b/gptqmodel/models/loader.py @@ -68,6 +68,7 @@ from ..utils.marlin import _marlin_capability_supported, _validate_marlin_device_support from ..utils.swordfish import _validate_swordfish_device_support from ..utils.model import ( + apply_no_placement_to_device_map, auto_dtype, convert_gptq_v1_to_v2_format, find_config_seq_len, @@ -80,6 +81,7 @@ is_embeddings_module_quantized, load_checkpoint_in_model_then_tie_weights, make_quant, + no_placement_module_names, simple_dispatch_model, ) from ._const import DEVICE, HAS_NPU, normalize_device @@ -1638,6 +1640,17 @@ def assign(mod, device_id): else: device_map = dict(explicit_device_map) log.info(f"Loader: honoring explicit device_map request: {device_map}") + original_device_map = dict(device_map) + # Checkpoint loading needs a non-overlapping map: parent and child entries + # would otherwise make Accelerate read the same PLE tensor on both devices. + device_map = apply_no_placement_to_device_map(model, device_map) + if device_map != original_device_map: + cpu_modules = sorted(no_placement_module_names(model)) + log.info(f"Loader: keeping no-placement modules on CPU: {cpu_modules}") + # Runtime dispatch keeps the parent entry so layer inputs still move to + # the right GPU, while the explicit CPU leaf blocks recursive PLE moves. + dispatch_device_map = dict(original_device_map) + dispatch_device_map.update(dict.fromkeys(no_placement_module_names(model), "cpu")) log.info(f"Loader: device_map = {device_map}") load_checkpoint_in_model = native_gguf_qspec is None @@ -1762,7 +1775,7 @@ def assign(mod, device_id): ) if native_gguf_qspec is not None: - model = simple_dispatch_model(model, device_map) + model = simple_dispatch_model(model, dispatch_device_map) _load_quantized_gguf_checkpoint_into_model( model=model, gguf_checkpoint_path=gguf_checkpoint_path, @@ -1770,7 +1783,7 @@ def assign(mod, device_id): ) else: # TODO: Why are we using this custom function and not dispatch_model? - model = simple_dispatch_model(model, device_map) + model = simple_dispatch_model(model, dispatch_device_map) if format_code == FORMAT.EXL3: qlinear_kernel = ExllamaV3TorchLinear if backend == BACKEND.EXL3_TORCH else ExllamaV3Linear diff --git a/gptqmodel/utils/model.py b/gptqmodel/utils/model.py index 9cf43d89f..ffeaf0d62 100644 --- a/gptqmodel/utils/model.py +++ b/gptqmodel/utils/model.py @@ -1223,6 +1223,90 @@ def wrapper(name): log.info("Model packed.") return quant_linear_cls + +def no_placement_module_names(model: nn.Module) -> set[str]: + """Resolve Transformers no-placement parameter patterns to leaf modules.""" + + patterns = getattr(model, "_no_placement_params", ()) or () + patterns = tuple(pattern for pattern in patterns if isinstance(pattern, str) and pattern) + if not patterns: + return set() + + names = set() + tensors = (*model.named_parameters(), *model.named_buffers()) + for tensor_name, _ in tensors: + if any(tensor_name == pattern or tensor_name.endswith(f".{pattern}") for pattern in patterns): + names.add(tensor_name.rsplit(".", 1)[0]) + return names + + +def _remove_redundant_device_map_children(device_map: Dict[str, Union[str, int]]) -> None: + """Remove child entries already covered by a same-device parent.""" + + for name in sorted(device_map, key=lambda item: item.count(".")): + if name not in device_map: + continue + prefix = f"{name}." if name else "" + for child_name in list(device_map): + if child_name != name and child_name.startswith(prefix) and device_map[child_name] == device_map[name]: + device_map.pop(child_name) + + +def _split_device_map_around_module( + model: nn.Module, + device_map: Dict[str, Union[str, int]], + ancestor_name: str, + target_name: str, + ancestor_device: Union[str, int], +) -> None: + """Replace one parent mapping with non-overlapping branches around a target.""" + + current_name = ancestor_name + while current_name != target_name: + current_module = model if not current_name else model.get_submodule(current_name) + relative_target = target_name if not current_name else target_name[len(current_name) + 1:] + path_child = relative_target.split(".", 1)[0] + + # Preserve direct tensors and sibling branches on the parent's device; + # descend only through the branch containing the excluded module. + for param_name, _ in (*current_module.named_parameters(recurse=False), *current_module.named_buffers(recurse=False)): + full_name = f"{current_name}.{param_name}" if current_name else param_name + device_map.setdefault(full_name, ancestor_device) + for child_name, _ in current_module.named_children(): + full_name = f"{current_name}.{child_name}" if current_name else child_name + if child_name != path_child: + device_map.setdefault(full_name, ancestor_device) + + current_name = f"{current_name}.{path_child}" if current_name else path_child + + +def apply_no_placement_to_device_map(model: nn.Module, device_map: Dict[str, Union[str, int]]) -> Dict[str, Union[str, int]]: + """Build a non-overlapping load map with excluded leaf modules on CPU.""" + + result = dict(device_map) + _remove_redundant_device_map_children(result) + for module_name in no_placement_module_names(model): + ancestors = [ + name + for name in result + if name != module_name and (not name or module_name.startswith(f"{name}.")) + ] + for ancestor_name in sorted(ancestors, key=lambda item: item.count("."), reverse=True): + ancestor_device = result.pop(ancestor_name) + _split_device_map_around_module( + model, + result, + ancestor_name, + module_name, + ancestor_device, + ) + # The checkpoint preloader expands parent and child entries independently; + # keep this map non-overlapping so the same tensor is not read on both devices. + result[module_name] = "cpu" + _remove_redundant_device_map_children(result) + return result + + def simple_dispatch_model(model, device_map): from accelerate.hooks import AlignDevicesHook, add_hook_to_module @@ -1249,7 +1333,15 @@ def simple_dispatch_model(model, device_map): else: main_device = [d for d in device_map.values() if d not in ["cpu", "disk"]][0] - cpu_offload_group = [(n, d) for n, d in device_map.items() if d == "cpu"] + # These modules perform their own CPU lookup and must remain resident; + # a normal CPU-offload hook would move their full weights back to the GPU. + resident_cpu_modules = no_placement_module_names(model) + module_names = dict(model.named_modules()) + cpu_offload_group = [ + (n, d) + for n, d in device_map.items() + if d == "cpu" and n not in resident_cpu_modules and n in module_names + ] prev_hook = None for idx, (n, d) in enumerate(cpu_offload_group): m = get_module_by_name_suffix(model, n) @@ -1261,10 +1353,18 @@ def simple_dispatch_model(model, device_map): for n, d in device_map.items(): if n == "": continue - m = get_module_by_name_suffix(model, n) + m = module_names.get(n) + if m is None: + # Fine-grained maps can contain direct parameter entries. + continue if d != "cpu": d = torch.device(d) - hook = AlignDevicesHook(d, io_same_device=True, place_submodules=True) + has_other_device_child = any( + child_name.startswith(f"{n}.") and child_device != device_map[n] + for child_name, child_device in device_map.items() + ) + # A mixed-device child means the parent hook may move inputs, but not descendants. + hook = AlignDevicesHook(d, io_same_device=True, place_submodules=not has_other_device_child) add_hook_to_module(m, hook) accelerate.utils.modeling.retie_parameters(model, tied_params) diff --git a/gptqmodel/utils/structure.py b/gptqmodel/utils/structure.py index 2d9a260b5..d7f5b2697 100644 --- a/gptqmodel/utils/structure.py +++ b/gptqmodel/utils/structure.py @@ -841,6 +841,11 @@ def __init__( alias_items = self._normalize_runtime_to_checkpoint_renamings(conversion_aliases) self._runtime_to_checkpoint_renamings = tuple(alias_items) self._runtime_to_checkpoint_converters = self._normalize_runtime_to_checkpoint_converters(conversion_aliases) + # Keep the same placement exclusions declared by the Transformers model. + no_placement_params = getattr(target_model, "_no_placement_params", ()) + self._no_placement_params = tuple( + pattern for pattern in (no_placement_params or ()) if isinstance(pattern, str) and pattern + ) self._lock = threading.RLock() @classmethod @@ -1264,11 +1269,14 @@ def reverse_hf_conversion_map(cls, conversion_mapping: Optional[Any]) -> Optiona return reversed_map or None @classmethod - def infer_hf_conversion_map_reversed(cls, *, target_model: Optional[nn.Module] = None) -> Optional[Any]: - if target_model is None: - return None - - model_type = getattr(getattr(target_model, "config", None), "model_type", None) + def infer_hf_conversion_map_reversed( + cls, + *, + target_model: Optional[nn.Module] = None, + model_type: Optional[str] = None, + ) -> Optional[Any]: + if model_type is None: + model_type = getattr(getattr(target_model, "config", None), "model_type", None) if isinstance(model_type, str): # Prefer the public transformers conversion registry and fall back to # older per-model mappings when needed. @@ -1287,6 +1295,8 @@ def infer_hf_conversion_map_reversed(cls, *, target_model: Optional[nn.Module] = if reversed_map is not None: return reversed_map + if target_model is None: + return None return cls.reverse_hf_conversion_map(getattr(target_model, "_checkpoint_conversion_mapping", None)) @staticmethod @@ -1961,15 +1971,40 @@ def _resolve_concat_checkpoint_tensor_sources( for converter in self._runtime_to_checkpoint_converters: if "Concatenate" not in converter.operation_names: continue - if len(converter.source_patterns) != 1 or len(converter.target_patterns) < 2: + if len(converter.source_patterns) != 1: continue runtime_pattern = converter.source_patterns[0] if _LazyWeightRenaming(runtime_pattern, runtime_pattern).rename_source_key(combined_name)[1] is None: continue + concat_operation = next( + operation + for operation in converter.operations + if type(operation).__name__ == "Concatenate" + ) + checkpoint_patterns = list(converter.target_patterns) + num_shards_attribute = getattr(concat_operation, "num_shards_attribute", None) + if len(checkpoint_patterns) == 1 and "*" in checkpoint_patterns[0] and num_shards_attribute: + # Qwen4-Exp stores one PLE embedding as config-counted shard names. + text_config = self.config + get_text_config = getattr(text_config, "get_text_config", None) + if callable(get_text_config): + text_config = get_text_config() + else: + text_config = getattr(text_config, "text_config", text_config) + num_shards = getattr(text_config, num_shards_attribute, None) + if not isinstance(num_shards, int) or num_shards <= 0: + continue + checkpoint_patterns = [ + checkpoint_patterns[0].replace("*", str(index)) + for index in range(num_shards) + ] + elif len(checkpoint_patterns) < 2: + continue + checkpoint_names = [] - for checkpoint_pattern in converter.target_patterns: + for checkpoint_pattern in checkpoint_patterns: renamed, matched_pattern = _LazyWeightRenaming( runtime_pattern, checkpoint_pattern, @@ -1988,18 +2023,35 @@ def _resolve_concat_checkpoint_tensor_sources( break checkpoint_names.append(resolved_name) - if len(checkpoint_names) != len(converter.target_patterns): + if len(checkpoint_names) != len(checkpoint_patterns): continue - concat_dim = 0 - for operation in converter.operations: - if type(operation).__name__ == "Concatenate": - concat_dim = getattr(operation, "dim", 0) - break + concat_dim = getattr(concat_operation, "dim", 0) return checkpoint_names, concat_dim return None + def _materialization_device_for_tensor( + self, + module_path: str, + rel_name: str, + default_device: torch.device, + ) -> torch.device: + """Prevent layer materialization from moving an excluded tensor off CPU.""" + + if self.is_no_placement_tensor(module_path, rel_name): + return torch.device("cpu") + return torch.device(default_device) + + def is_no_placement_tensor(self, module_path: str, rel_name: str) -> bool: + """Check Transformers' accelerator placement exclusions.""" + + full_name = self._join_tensor_name(module_path, rel_name) + for pattern in self._no_placement_params: + if full_name == pattern or full_name.endswith(f".{pattern}"): + return True + return False + def _resolve_direct_checkpoint_tensor_source( self, module_path: str, @@ -2446,6 +2498,7 @@ def _copy_checkpoint_tensors_into_submodule( ) if kind == "param": + tensor_device = self._materialization_device_for_tensor(module_path, rel_name, device) target_param = t_params.get(rel_name) if target_param is None: raise RuntimeError( @@ -2459,7 +2512,7 @@ def _copy_checkpoint_tensors_into_submodule( source_shape=tuple(tensor.shape), ) ) - target_param_new = _ensure_target_storage_on_device_(target_param, device) + target_param_new = _ensure_target_storage_on_device_(target_param, tensor_device) if target_param_new is not target_param: t_parent, leaf = _get_parent_and_leaf_by_path(target_submodule, rel_name) setattr(t_parent, leaf, target_param_new) @@ -2535,6 +2588,7 @@ def _copy_checkpoint_tensors_into_submodule( split_dim=split_dim, )) if kind == "param": + tensor_device = self._materialization_device_for_tensor(module_path, rel_name, device) target_param = t_params.get(rel_name) if target_param is None: raise RuntimeError(self._materialization_issue_message( @@ -2563,7 +2617,7 @@ def _copy_checkpoint_tensors_into_submodule( split_index=split_index, split_dim=split_dim, )) - target_param_new = _ensure_target_storage_on_device_(target_param, device) + target_param_new = _ensure_target_storage_on_device_(target_param, tensor_device) if target_param_new is not target_param: t_parent, leaf = _get_parent_and_leaf_by_path(target_submodule, rel_name) setattr(t_parent, leaf, target_param_new) diff --git a/tests/models/test_qwen4_exp.py b/tests/models/test_qwen4_exp.py new file mode 100644 index 000000000..a653a7de9 --- /dev/null +++ b/tests/models/test_qwen4_exp.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: 2026 ModelCloud.ai +# SPDX-License-Identifier: Apache-2.0 + +from model_test import ModelTest + +from gptqmodel.quantization.config import GcMode + + +class TestQwen3_8FlashNext(ModelTest): + NATIVE_MODEL_ID = "/monster/data/model/Qwen3.8-Flash-Next" + TRUST_REMOTE_CODE = False + USE_FLASH_ATTN = False + EVAL_BATCH_SIZE = 16 + EVAL_SINGLE_GPU = False + + EVAL_TASKS_SLOW = { + "arc_challenge": { + "acc": {"value": 0.6194539249146758, "floor_pct": 0.04}, + "acc_norm": {"value": 0.6100682593856656, "floor_pct": 0.04}, + }, + } + EVAL_TASKS_FAST = ModelTest.derive_fast_eval_tasks(EVAL_TASKS_SLOW) + + MODEL_COMPAT_FAST_LAYER_POSITION = "first" + SAVE_PATH = "./temp/qwen4_exp_test" + + def _build_quantize_config(self): + config = super()._build_quantize_config() + # Drain 1,500+ expert pack jobs before replaying the next layer. + config.wait_for_submodule_finalizers = True + # Release temporary replay buffers after each stage. + config.gc_mode = GcMode.ON_STAGE_END + return config + + def test_qwen3_8_flash_next(self): + self.quantize_and_evaluate() + + +__all__ = ["TestQwen3_8FlashNext"] diff --git a/tests/test_lazy_turtle_conversion_mapping.py b/tests/test_lazy_turtle_conversion_mapping.py index 819a246b8..750f88ec0 100644 --- a/tests/test_lazy_turtle_conversion_mapping.py +++ b/tests/test_lazy_turtle_conversion_mapping.py @@ -10,6 +10,7 @@ from safetensors.torch import save_file from torch import nn +from gptqmodel.models.base import BaseQModel from gptqmodel.models.definitions.deepseek_ocr2 import DeepSeekOCR2QModel from gptqmodel.models.definitions.deepseek_v4 import DeepSeekV4QModel from gptqmodel.models.definitions.gemma3 import Gemma3ForConditionalGenerationGPTQ @@ -34,6 +35,7 @@ def _build_lazy_turtle( tmp_path: Path, checkpoint_tensors: dict[str, torch.Tensor], *, + config=None, module_tree=None, hf_conversion_map_reversed=None, target_model: nn.Module | None = None, @@ -45,7 +47,7 @@ def _build_lazy_turtle( _write_checkpoint_index(model_dir, shard_name, checkpoint_tensors) turtle = LazyTurtle.maybe_create( model_local_path=str(model_dir), - config=SimpleNamespace(_experts_implementation=None), + config=config or SimpleNamespace(_experts_implementation=None), model_init_kwargs={"device_map": {"": "cpu"}}, module_tree=module_tree, hf_conversion_map_reversed=hf_conversion_map_reversed, @@ -166,6 +168,24 @@ def __init__(self): self.model.language_model.layers = nn.ModuleList([_FusedDenseLayerShell()]) +class _Qwen4NgramShell(nn.Module): + _no_placement_params = ["ple.ple_embedding.ngram_embedding.weight"] + + def __init__(self): + super().__init__() + self.config = SimpleNamespace( + model_type="qwen4_exp", + text_config=SimpleNamespace(split_ngram_parts=3), + ) + self.model = nn.Module() + self.model.language_model = nn.Module() + layer = nn.Module() + layer.ple = nn.Module() + layer.ple.ple_embedding = nn.Module() + layer.ple.ple_embedding.ngram_embedding = nn.Embedding(6, 2, device="meta") + self.model.language_model.layers = nn.ModuleList([layer]) + + class _MiniMaxM3SharedExpertsMlpShell(nn.Module): def __init__(self, hidden_dim: int = 4, intermediate_dim: int = 3): super().__init__() @@ -255,8 +275,9 @@ def __init__(self, dim: int = 0): class Concatenate: - def __init__(self, dim: int = 0): + def __init__(self, dim: int = 0, num_shards_attribute: str | None = None): self.dim = dim + self.num_shards_attribute = num_shards_attribute class ErnieFuseAndSplitTextVisionExperts: @@ -1081,6 +1102,56 @@ def test_lazy_turtle_materializes_fused_dense_mlp_from_split_gate_up_checkpoint( assert torch.equal(weight, torch.cat([gate, up], dim=0)) +def test_lazy_turtle_materializes_config_counted_shards_on_cpu_for_no_placement_param(tmp_path): + reversed_map = LazyTurtle.reverse_hf_conversion_map( + [ + _WeightConverterStub( + source_patterns="ngram_embedding.shard_*.weight", + target_patterns="ngram_embedding.weight", + operations=[Concatenate(dim=0, num_shards_attribute="split_ngram_parts")], + ), + ] + ) + assert reversed_map is not None + + shards = [ + torch.arange(0, 4, dtype=torch.float32).reshape(2, 2), + torch.arange(4, 6, dtype=torch.float32).reshape(1, 2), + torch.arange(6, 12, dtype=torch.float32).reshape(3, 2), + ] + checkpoint_tensors = { + f"model.language_model.layers.0.ple.ple_embedding.ngram_embedding.shard_{index}.weight": shard + for index, shard in enumerate(shards) + } + shell = _Qwen4NgramShell() + turtle = _build_lazy_turtle( + tmp_path, + checkpoint_tensors, + config=shell.config, + hf_conversion_map_reversed=reversed_map, + target_model=shell, + ) + + ngram_embedding = shell.model.language_model.layers[0].ple.ple_embedding.ngram_embedding + turtle.materialize_submodule( + target_model=shell, + target_submodule=ngram_embedding, + device=torch.device("cuda:0"), + module_path="model.language_model.layers.0.ple.ple_embedding.ngram_embedding", + show_progress=False, + ) + + assert ngram_embedding.weight.device.type == "cpu" + assert torch.equal(ngram_embedding.weight, torch.cat(shards, dim=0)) + + qmodel = BaseQModel.__new__(BaseQModel) + nn.Module.__init__(qmodel) + qmodel.model = shell + qmodel.turtle_model = turtle + assert qmodel.has_forward_device_overrides() is True + assert qmodel.forward_device_for_module(ngram_embedding, torch.device("cuda:0")) == torch.device("cpu") + + def test_lazy_turtle_sync_all_meta_materializes_fused_dense_mlp_from_split_gate_up_checkpoint(tmp_path): reversed_map = LazyTurtle.reverse_hf_conversion_map( [ diff --git a/tests/test_qwen4_exp_support.py b/tests/test_qwen4_exp_support.py new file mode 100644 index 000000000..5adf3aaf4 --- /dev/null +++ b/tests/test_qwen4_exp_support.py @@ -0,0 +1,338 @@ +# SPDX-FileCopyrightText: 2026 ModelCloud.ai +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from gptqmodel.models import auto +from gptqmodel.models.definitions.qwen4_exp import Qwen4ExpQModel +from gptqmodel.models.loader import _convert_model_with_defuser +from gptqmodel.utils.model import apply_no_placement_to_device_map, simple_dispatch_model +from gptqmodel.utils.structure import LazyTurtle + + +def _outer_config(num_experts=3): + return SimpleNamespace(text_config=SimpleNamespace(num_experts=num_experts)) + + +class _TinyNoPlacementModel(nn.Module): + _no_placement_params = ["ple.ple_embedding.ngram_embedding.weight"] + + def __init__(self): + super().__init__() + self.model = nn.Module() + self.model.language_model = nn.Module() + self.model.language_model.layers = nn.ModuleList([nn.Module(), nn.Module()]) + layer = self.model.language_model.layers[1] + layer.proj = nn.Linear(4, 4) + layer.ple = nn.Module() + layer.ple.ple_embedding = nn.Module() + layer.ple.ple_embedding.ngram_embedding = nn.Embedding(8, 4) + layer.ple.key_proj = nn.Linear(4, 4) + + +def test_qwen4_exp_model_type_selects_definition(monkeypatch): + fake_config = SimpleNamespace(model_type="qwen4_exp") + monkeypatch.setattr(auto, "resolve_trust_remote_code", lambda path, trust_remote_code=False: trust_remote_code) + monkeypatch.setattr(auto.AutoConfig, "from_pretrained", lambda *args, **kwargs: fake_config) + + assert auto.check_and_get_model_definition("/tmp/qwen3.8-flash-next") is Qwen4ExpQModel + + +def test_qwen4_exp_quantized_load_keeps_ple_embedding_on_cpu(): + model = _TinyNoPlacementModel() + layer_name = "model.language_model.layers.1" + embedding_name = f"{layer_name}.ple.ple_embedding.ngram_embedding" + + device_map = apply_no_placement_to_device_map(model, {layer_name: "cuda:1"}) + + assert layer_name not in device_map + assert device_map[f"{layer_name}.proj"] == "cuda:1" + assert device_map[f"{layer_name}.ple.key_proj"] == "cuda:1" + assert device_map[embedding_name] == "cpu" + assert not any( + embedding_name.startswith(f"{name}.") + for name in device_map + if name != embedding_name + ) + + +def test_qwen4_exp_dispatch_does_not_move_cpu_ple_under_gpu_parent(monkeypatch): + import accelerate + + model = _TinyNoPlacementModel() + layer_name = "model.language_model.layers.1" + embedding_name = f"{layer_name}.ple.ple_embedding.ngram_embedding" + device_map = {layer_name: "cuda:1", embedding_name: "cpu"} + added_hooks = [] + + monkeypatch.setattr(accelerate.utils.modeling, "find_tied_parameters", lambda model: []) + monkeypatch.setattr(accelerate.utils.modeling, "retie_parameters", lambda model, tied: None) + monkeypatch.setattr( + accelerate, + "cpu_offload_with_hook", + lambda *args, **kwargs: pytest.fail("no-placement PLE must remain resident on CPU"), + ) + monkeypatch.setattr( + accelerate.hooks, + "add_hook_to_module", + lambda module, hook: added_hooks.append((module, hook)), + ) + + simple_dispatch_model(model, device_map) + + assert len(added_hooks) == 1 + assert added_hooks[0][0] is model.model.language_model.layers[1] + assert added_hooks[0][1].place_submodules is False + + +def test_qwen4_exp_module_tree_quantizes_selected_attention_and_mlp_linears(): + modules = Qwen4ExpQModel.simple_layer_modules( + _outer_config(), + SimpleNamespace(dynamic=None), + ) + flat = {name for block in modules for name in block} + captured = { + name + for block in Qwen4ExpQModel.full_layer_modules( + _outer_config(), + include_capture_only=True, + ) + for name in block + } + + assert Qwen4ExpQModel.extract_layers_node() == ["model.language_model.layers"] + assert Qwen4ExpQModel.pre_lm_head_norm_module == "model.language_model.hyper_connection_mixer" + assert Qwen4ExpQModel.out_of_model_tensors == {"prefixes": ["mtp"]} + assert Qwen4ExpQModel.hf_conversion_model_type_alias == "qwen4_exp_text" + + assert modules[:7] == [ + ["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj"], + ["self_attn.o_proj"], + ["linear_attn.in_proj_qkv"], + ["linear_attn.in_proj_z"], + ["linear_attn.out_proj"], + ["mlp.shared_expert.gate_proj", "mlp.shared_expert.up_proj"], + ["mlp.shared_expert.down_proj"], + ] + assert len(modules) == 9 + assert flat == { + "self_attn.q_proj", + "self_attn.k_proj", + "self_attn.v_proj", + "self_attn.o_proj", + "linear_attn.in_proj_qkv", + "linear_attn.in_proj_z", + "linear_attn.out_proj", + "mlp.shared_expert.gate_proj", + "mlp.shared_expert.up_proj", + "mlp.shared_expert.down_proj", + "mlp.experts.0.gate_proj", + "mlp.experts.0.up_proj", + "mlp.experts.0.down_proj", + "mlp.experts.1.gate_proj", + "mlp.experts.1.up_proj", + "mlp.experts.1.down_proj", + "mlp.experts.2.gate_proj", + "mlp.experts.2.up_proj", + "mlp.experts.2.down_proj", + } + + for name in ( + "self_attn.indexer.index_qk_proj:!", + "self_attn.indexer.q_layernorm:!", + "linear_attn.conv1d:!", + "linear_attn.in_proj_b:!", + "linear_attn.in_proj_a:!", + "linear_attn.norm:!", + "mlp.gate:!", + "mlp.shared_expert_gate:!", + ): + assert name in captured + + shared = [name for block in modules for name in block if name.startswith("mlp.shared_expert.")] + assert shared == [ + "mlp.shared_expert.gate_proj", + "mlp.shared_expert.up_proj", + "mlp.shared_expert.down_proj", + ] + + +def _tiny_qwen4_exp_text_model(): + transformers = pytest.importorskip("transformers") + config_cls = getattr(transformers, "Qwen4ExpTextConfig", None) + model_cls = getattr(transformers, "Qwen4ExpForCausalLM", None) + if config_cls is None or model_cls is None: + pytest.skip("Installed Transformers does not provide Qwen4Exp") + + config = config_cls( + vocab_size=64, + hidden_size=16, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=4, + moe_intermediate_size=8, + shared_expert_intermediate_size=8, + num_experts=3, + num_experts_per_tok=2, + max_position_embeddings=32, + layer_types=["linear_attention", "qwen_sparse_attention"], + linear_num_key_heads=2, + linear_num_value_heads=4, + linear_key_head_dim=4, + linear_value_head_dim=4, + linear_conv_kernel_dim=2, + hc_count=2, + hc_lowrank=4, + ple_layer_ids=[], + bos_token_id=1, + eos_token_id=2, + pad_token_id=0, + indexer_n_heads=2, + indexer_kv_heads=1, + indexer_head_dim=4, + indexer_budget=8, + indexer_compress_ratio=2, + ) + return model_cls(config).eval() + + +def test_qwen4_exp_outer_model_reuses_text_checkpoint_conversion_map(): + transformers = pytest.importorskip("transformers") + if getattr(transformers, "Qwen4ExpConfig", None) is None: + pytest.skip("Installed Transformers does not provide Qwen4Exp") + + outer_model = SimpleNamespace(config=SimpleNamespace(model_type="qwen4_exp")) + reversed_map = Qwen4ExpQModel.resolve_hf_conversion_map_reversed(target_model=outer_model) + + assert reversed_map is not None + assert any( + "ngram_embedding.weight" in converter.source_patterns + and "ngram_embedding.shard_*.weight" in converter.target_patterns + for converter in reversed_map + if hasattr(converter, "source_patterns") and hasattr(converter, "target_patterns") + ) + + +def test_qwen4_exp_outer_lazy_materialization_loads_ple_shards(tmp_path): + transformers = pytest.importorskip("transformers") + config_cls = getattr(transformers, "Qwen4ExpConfig", None) + model_cls = getattr(transformers, "Qwen4ExpForConditionalGeneration", None) + if config_cls is None or model_cls is None: + pytest.skip("Installed Transformers does not provide Qwen4Exp") + + text_config = { + "vocab_size": 64, + "hidden_size": 32, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 8, + "moe_intermediate_size": 16, + "shared_expert_intermediate_size": 16, + "num_experts": 3, + "num_experts_per_tok": 2, + "max_position_embeddings": 32, + "layer_types": ["qwen_sparse_attention", "linear_attention"], + "linear_num_key_heads": 2, + "linear_num_value_heads": 4, + "linear_key_head_dim": 8, + "linear_value_head_dim": 8, + "linear_conv_kernel_dim": 2, + "hc_count": 2, + "hc_lowrank": 8, + "ple_layer_ids": [2], + "ple_embed_dim": 16, + "ple_conv_kernel_size": 2, + "ngram_size": 3, + "heads_per_ngram": 2, + "ngram_vocab_size_base": 32, + "make_ngram_vocab_size_divisible_by": 2, + "split_ngram_parts": 2, + "indexer_n_heads": 2, + "indexer_kv_heads": 1, + "indexer_head_dim": 8, + "indexer_budget": 8, + "indexer_compress_ratio": 2, + "bos_token_id": 1, + "eos_token_id": 2, + "pad_token_id": 0, + } + vision_config = { + "depth": 1, + "hidden_size": 32, + "intermediate_size": 64, + "num_heads": 4, + "in_channels": 3, + "patch_size": 2, + "spatial_merge_size": 1, + "temporal_patch_size": 1, + "out_hidden_size": 32, + "num_position_embeddings": 16, + } + config = config_cls( + text_config=text_config, + vision_config=vision_config, + image_token_id=60, + video_token_id=61, + vision_start_token_id=62, + vision_end_token_id=63, + ) + source = model_cls(config).eval() + expected = source.model.language_model.layers[1].ple.ple_embedding.ngram_embedding.weight.detach().clone() + source.save_pretrained(tmp_path) + + with torch.device("meta"): + shell = model_cls(config).eval() + assert _convert_model_with_defuser(Qwen4ExpQModel, shell, cleanup_original=False) + turtle = LazyTurtle.maybe_create( + model_local_path=str(tmp_path), + config=shell.config, + model_init_kwargs={"device_map": {"": "cpu"}}, + module_tree=Qwen4ExpQModel.module_tree, + hf_conversion_map_reversed=Qwen4ExpQModel.resolve_hf_conversion_map_reversed(target_model=shell), + target_model=shell, + ) + assert turtle is not None + + layer = shell.model.language_model.layers[1] + turtle.materialize_submodule( + target_model=shell, + target_submodule=layer, + device=torch.device("cpu"), + module_path="model.language_model.layers.1", + show_progress=False, + ) + + actual = layer.ple.ple_embedding.ngram_embedding.weight + assert actual.device.type == "cpu" + assert not any(parameter.is_meta for parameter in layer.parameters()) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_qwen4_exp_defuser_preserves_model_type_and_forward(): + model = _tiny_qwen4_exp_text_model() + input_ids = torch.tensor([[1, 7, 8, 2]]) + original_model_type = model.config.model_type + packed_experts = model.model.layers[0].mlp.experts + assert hasattr(packed_experts, "gate_up_proj") + + with torch.inference_mode(): + expected = model(input_ids=input_ids, use_cache=False).logits + + assert _convert_model_with_defuser(Qwen4ExpQModel, model, cleanup_original=False) is True + experts = model.model.layers[0].mlp.experts + + assert model.config.model_type == original_model_type + assert not hasattr(experts, "gate_up_proj") + assert isinstance(experts[0].gate_proj, torch.nn.Linear) + assert isinstance(experts[0].up_proj, torch.nn.Linear) + assert isinstance(experts[0].down_proj, torch.nn.Linear) + + with torch.inference_mode(): + actual = model(input_ids=input_ids, use_cache=False).logits + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-7) diff --git a/tests/test_subset_plan.py b/tests/test_subset_plan.py index 1054f1292..39b15b8b1 100644 --- a/tests/test_subset_plan.py +++ b/tests/test_subset_plan.py @@ -368,6 +368,71 @@ def _group_key(name: str): assert subset["self_attn.v_proj"].state["preferred_quant_device"] == torch.device("cuda:0") +def test_build_subset_plan_honors_model_forward_device_override(): + looper = _make_looper() + looper._quant_devices = [torch.device("cuda:0")] + looper._dense_quant_devices = [torch.device("cuda:0")] + looper._dense_vram_strategy_explicit = True + + processor = _StubProcessor(ExecutionConfig(require_fwd=True)) + subset = {"ple.key_proj": _make_named_module("ple.key_proj")} + ngram_embedding = torch.nn.Embedding(8, 4) + full = { + "ple.key_proj": subset["ple.key_proj"].module, + "ple.ple_embedding.ngram_embedding": ngram_embedding, + } + looper.gptq_model.forward_device_for_module = ( + lambda module, planned_device: torch.device("cpu") if module is ngram_embedding else planned_device + ) + looper.gptq_model.has_forward_device_overrides = lambda: True + + plan = build_subset_plan( + looper, + processor=processor, + subset=subset, + subset_index=0, + subset_total=1, + full=full, + fallback=True, + layer_inputs=[[torch.zeros(1, 4)]], + planning_layer_modules=_planning_blocks(("ple.ple_embedding:!", "ple.key_proj")), + ) + + assert plan.forward_device_map["ple.key_proj"] == torch.device("cuda:0") + assert plan.forward_device_map["ple.ple_embedding.ngram_embedding"] == torch.device("cpu") + + +def test_build_subset_plan_activates_model_override_without_explicit_vram_strategy(): + looper = _make_looper() + processor = _StubProcessor(ExecutionConfig(require_fwd=True)) + subset = {"mlp.experts.0.gate_proj": _make_named_module("mlp.experts.0.gate_proj")} + ngram_embedding = torch.nn.Embedding(8, 4) + full = { + "mlp.experts.0.gate_proj": subset["mlp.experts.0.gate_proj"].module, + "ple.ple_embedding.ngram_embedding": ngram_embedding, + } + looper.gptq_model.forward_device_for_module = ( + lambda module, planned_device: torch.device("cpu") if module is ngram_embedding else planned_device + ) + looper.gptq_model.has_forward_device_overrides = lambda: True + + plan = build_subset_plan( + looper, + processor=processor, + subset=subset, + subset_index=0, + subset_total=2, + full=full, + fallback=True, + layer_inputs=[[torch.zeros(1, 4)]], + planning_layer_modules=_planning_blocks(("mlp.experts.0.gate_proj",)), + ) + + assert plan.forward_device_map["ple.ple_embedding.ngram_embedding"] == torch.device("cpu") + assert plan.forward_mode == "serial" + assert plan.restore_forward_device_overrides is False + + def test_build_subset_plan_split_pools_pin_dense_subset_and_balance_experts(): looper = _make_looper() looper._quant_devices = [