# Copyright (C) 2025 the HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # 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 __future__ import annotations from collections import defaultdict from copy import deepcopy from typing import TYPE_CHECKING from .core_model_loading import ( Chunk, Concatenate, ErnieFuseAndSplitTextVisionExperts, MergeModulelist, PrefixChange, Transpose, WeightConverter, WeightRenaming, WeightTransform, ) if TYPE_CHECKING: from .modeling_utils import PreTrainedModel from .quantizers import HfQuantizer _MODEL_TO_CONVERSION_PATTERN = { # Mixtral-style MoE "minimax": "mixtral", "minimax_m2": "mixtral", # Qwen2-style MoE "afmoe": "qwen2_moe", "deepseek_v2": "qwen2_moe", "deepseek_v3": "qwen2_moe", "dots1": "qwen2_moe", "ernie4_5_moe": "qwen2_moe", "glm4_moe": "qwen2_moe", "glm4_moe_lite": "qwen2_moe", "glm_moe_dsa": "qwen2_moe", "glm4v_moe": "qwen2_moe", "longcat_flash": "qwen2_moe", "solar_open": "qwen2_moe", "qwen3_moe": "qwen2_moe", "qwen3_omni_moe": "qwen2_moe", "qwen3_omni_moe_thinker": "qwen2_moe", "qwen3_next": "qwen2_moe", "hunyuan_v1_moe": "qwen2_moe", "flex_olmo": "qwen2_moe", "olmoe": "qwen2_moe", "exaone_moe": "qwen2_moe", "cohere2_moe": "qwen2_moe", "rt_detr_v2": "rt_detr", "pp_doclayout_v2": "rt_detr", "pp_doclayout_v3": "rt_detr", "sam3_tracker_video": "sam3_tracker", "altclip_vision_model": "clip_vision_model", "chinese_clip_vision_model": "clip_vision_model", "clipseg_vision_model": "clip_vision_model", "metaclip_2_vision_model": "clip_vision_model", "mlcd_vision": "clip_vision_model", "mlcd": "clip_vision_model", "siglip_vision_model": "clip_vision_model", "siglip2_vision_model": "clip_vision_model", "xclip_vision_model": "clip_vision_model", "clipseg_text_model": "clip_text_model", "metaclip_2_text_model": "clip_text_model", "siglip_text_model": "clip_text_model", "siglip2_text_model": "clip_text_model", "xclip_text_model": "clip_text_model", "paligemma": "llava", "aya_vision": "llava", "got_ocr2": "llava", "gemma3": "llava", "internvl": "llava", "vipllava": "llava", "mistral3": "llava", "pp_chart2table": "llava", "llava_next_video": "llava_next", "llava_onevision": "llava_next", # class-based mappings "PaliGemmaModel": "LlavaModel", "AyaVisionModel": "LlavaModel", "GotOcr2Model": "LlavaModel", "Gemma3Model": "LlavaModel", "InternVLModel": "LlavaModel", "VipLlavaModel": "LlavaModel", "Mistral3Model": "LlavaModel", "PPChart2TableModel": "LlavaModel", "LlavaNextModel": "LlavaModel", "LlavaNextVideoModel": "LlavaModel", "LlavaOnevisionModel": "LlavaModel", "FuyuModel": "LlavaModel", "MllamaModel": "LlavaModel", "MaskFormerDetrDecoder": "DetrModel", "Qwen2_5_VLForConditionalGeneration": "Qwen2VLForConditionalGeneration", # ViT-style vision models (old HuggingFace checkpoint format → new modular format) "ASTModel": "ViTModel", "BeitModel": "ViTModel", "DeiTModel": "ViTModel", "IJepaModel": "ViTModel", "ViTMAEModel": "ViTModel", "ViTMSNModel": "ViTModel", "VivitModel": "ViTModel", } def _build_checkpoint_conversion_mapping(): mapping = { "hrm_text": [ WeightConverter( source_patterns="mlp.gate_up_proj.weight", target_patterns=["mlp.gate_proj.weight", "mlp.up_proj.weight"], operations=[Chunk(dim=0)], ), WeightConverter( source_patterns="attn.gqkv_proj.weight", target_patterns=[ "self_attn.gate_proj.weight", "self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight", ], operations=[Chunk(dim=0)], ), WeightRenaming(source_patterns=r"\.attn\.o_proj\.", target_patterns=".self_attn.o_proj."), ], "ViTModel": [ WeightRenaming(r"encoder\.layer\.", "layers."), WeightRenaming("attention.query", "q_proj"), WeightRenaming("attention.key", "k_proj"), WeightRenaming("attention.value", "v_proj"), WeightRenaming("attention.output.dense", "attention.o_proj"), WeightRenaming("intermediate.dense", "mlp.fc1"), WeightRenaming("output.dense", "mlp.fc2"), ], "ViTMSNForImageClassification": [ WeightRenaming(r"^encoder\.", "vit.encoder."), ], "ViTMAEForPreTraining": [ WeightRenaming("attention.query", "q_proj"), WeightRenaming("attention.key", "k_proj"), WeightRenaming("attention.value", "v_proj"), WeightRenaming("attention.output.dense", "attention.o_proj"), WeightRenaming("intermediate.dense", "mlp.fc1"), WeightRenaming("output.dense", "mlp.fc2"), ], "BeitBackbone": [ WeightRenaming(r"^fpn1\.0\.", "fpn.fpn1.conv_transpose1."), WeightRenaming(r"^fpn1\.1\.", "fpn.fpn1.normalization."), WeightRenaming(r"^fpn1\.3\.", "fpn.fpn1.conv_transpose2."), WeightRenaming(r"^fpn2\.0\.", "fpn.fpn2."), WeightRenaming(r"^encoder\.layer\.", "beit.encoder.layer."), WeightRenaming(r"^embeddings\.", "beit.embeddings."), WeightRenaming("attention.query", "q_proj"), WeightRenaming("attention.key", "k_proj"), WeightRenaming("attention.value", "v_proj"), WeightRenaming("attention.output.dense", "attention.o_proj"), WeightRenaming("intermediate.dense", "mlp.fc1"), WeightRenaming("output.dense", "mlp.fc2"), ], "BeitForSemanticSegmentation": [ WeightRenaming(r"(? we end up with 2 tensors Concatenate(dim=1), # each process has 2 tensors, gate and up, we concat them into gate_up ], # we want the loading to add this shard operation here. Though we can't shard after concats and merge, needs to be first ), WeightConverter( source_patterns=[ ".experts.*.w2.weight", ], target_patterns=".experts.down_proj", # target key gets the list of two tensors operations=[ MergeModulelist( dim=0 ), # each process has two lists of tensors, we cat each list. -> we end up with 2 tensors ], # we want the loading to add this shard operation here. Though we can't shard after concats and merge, needs to be first ), ], "qwen2_moe": [ WeightConverter( source_patterns=[ "mlp.experts.*.gate_proj.weight", "mlp.experts.*.up_proj.weight", ], target_patterns="mlp.experts.gate_up_proj", operations=[MergeModulelist(dim=0), Concatenate(dim=1)], ), WeightConverter( source_patterns="mlp.experts.*.down_proj.weight", target_patterns="mlp.experts.down_proj", operations=[MergeModulelist(dim=0)], ), ], "qwen3_vl_moe": [ WeightConverter( source_patterns="mlp.experts.gate_up_proj", target_patterns="mlp.experts.gate_up_proj", operations=[Transpose(1, 2, check_dims=True)], ), WeightConverter( source_patterns="mlp.experts.down_proj", target_patterns="mlp.experts.down_proj", operations=[Transpose(1, 2, check_dims=True)], ), ], "phimoe": [ WeightRenaming(".block_sparse_moe.", ".mlp."), WeightRenaming(".gate.weight", ".router.weight"), WeightConverter( source_patterns=[ ".experts.*.w1.weight", ".experts.*.w3.weight", ], target_patterns=".experts.gate_up_proj", operations=[MergeModulelist(dim=0), Concatenate(dim=1)], ), WeightConverter( source_patterns=".experts.*.w2.weight", target_patterns=".experts.down_proj", operations=[MergeModulelist(dim=0)], ), ], "lfm2_moe": [ WeightConverter( source_patterns=[ "feed_forward.experts.*.w1.weight", "feed_forward.experts.*.w3.weight", ], target_patterns="feed_forward.experts.gate_up_proj", operations=[MergeModulelist(dim=0), Concatenate(dim=1)], ), WeightConverter( source_patterns="feed_forward.experts.*.w2.weight", target_patterns="feed_forward.experts.down_proj", operations=[MergeModulelist(dim=0)], ), ], "ernie4_5_vl_moe": [ # vision WeightRenaming("vision_model", "vision_tower"), # resampler WeightRenaming("spatial_linear.0", "spatial_linear.fc1"), WeightRenaming("spatial_linear.2", "spatial_linear.fc2"), WeightRenaming("spatial_linear.3", "spatial_linear.ln"), WeightRenaming("temporal_linear.0", "temporal_linear.fc1"), WeightRenaming("temporal_linear.2", "temporal_linear.fc2"), WeightRenaming("temporal_linear.3", "temporal_linear.ln"), # language model WeightRenaming(r"(? None: """ Register a conversion mapping for a model type string or a class name. Class names take priority over `model_type` strings during lookup (see `extract_weight_conversions_for_model`), making it possible to define task-head-specific or class-specific conversions that differ from the shared `model_type` baseline. """ global _checkpoint_conversion_mapping_cache if _checkpoint_conversion_mapping_cache is None: _checkpoint_conversion_mapping_cache = _build_checkpoint_conversion_mapping() if model_type_or_class_name in _checkpoint_conversion_mapping_cache and not overwrite: raise ValueError( f"Conversion mapping for '{model_type_or_class_name}' already exists. Pass overwrite=True to replace it." ) _checkpoint_conversion_mapping_cache[model_type_or_class_name] = mapping def extract_weight_conversions_for_model( model: PreTrainedModel, ) -> list[WeightTransform] | None: """ Return the registered conversion list for `model`, or `None` if none exists. Looks up by class name first (enables task-head-specific overrides), then falls back to `model.config.model_type`. Transforms are returned unmodified; the caller sets `scope_prefix` on each transform for sub-module isolation. """ class_name = type(model).__name__ model_type = model.config.model_type # Class name takes priority — allows ForXxx-specific overrides conversions = get_checkpoint_conversion_mapping(class_name) if conversions is None and model_type: conversions = get_checkpoint_conversion_mapping(model_type) return conversions def get_model_conversion_mapping( model: PreTrainedModel, key_mapping: dict[str, str] | None = None, hf_quantizer: HfQuantizer | None = None, add_legacy: bool = True, ) -> list[WeightTransform]: """ Collect the ordered list of weight transforms for `model` (used during loading and, when reversed, during saving). Each `PreTrainedModel` sub-module is looked up by class name then `model_type`. Root transforms are applied globally; sub-module transforms have their `scope_prefix` set so they only match keys under that prefix. After any sub-module is processed, both its class name and `model_type` are marked seen to prevent `XForY` / `XModel` pairs from applying the same mapping twice via different lookup paths. """ from .modeling_utils import PreTrainedModel # note: this function is used in PEFT, so changing the API requires coordination weight_conversions = [] # Load models with explicit, user-provided key mapping if key_mapping is not None: weight_conversions = [WeightRenaming(source_patterns=k, target_patterns=v) for k, v in key_mapping.items()] # Maps each identifier (class name or model_type) to the module paths that have # already claimed it. A later module is skipped only when one of those paths is # an ancestor of the current module path — siblings are never ancestors of each # other, so two sibling sub-models with the same model_type both get their own # scoped transforms. A child is an ancestor of everything nested under it, which # prevents a parent's transforms from being duplicated with a scoped copy for the child. seen_identifiers: defaultdict[str, list[str]] = defaultdict(list) for module_name, submodule in model.named_modules(): # Skip if it's not a submodel if not isinstance(submodule, PreTrainedModel): continue class_name = type(submodule).__name__ model_type = submodule.config.model_type # Skip if an ancestor already claimed this class (its unscoped transforms already cover this subtree). if any(seen == "" or module_name.startswith(seen + ".") for seen in seen_identifiers[class_name]): continue # Class name takes priority — a class-specific mapping bypasses the model_type # deduplication check (e.g. LlavaModel nested inside LlavaForConditionalGeneration # must still get its own scoped mapping even after "llava" is marked seen). conversions = get_checkpoint_conversion_mapping(class_name) found_via_class = conversions is not None if not found_via_class: # Same ancestor check as above, but via model_type for modules without a class-specific mapping. if model_type and any( seen == "" or module_name.startswith(seen + ".") for seen in seen_identifiers[model_type] ): continue if model_type is not None: conversions = get_checkpoint_conversion_mapping(model_type) if conversions is None: continue is_root_model = module_name == "" if not is_root_model: # Scope each transform so it only matches keys under this sub-module's prefix. for transform in conversions: transform.scope_prefix = module_name weight_conversions.extend(conversions) seen_identifiers[class_name].append(module_name) # Only record model_type when the hit was via model_type. When the hit was via # class name, other sub-modules sharing the same model_type but without a # class-specific mapping (e.g. DetrModel under DetrForSegmentation) must still # be reachable so their base transforms are picked up and scoped. if not found_via_class and model_type: seen_identifiers[model_type].append(module_name) if add_legacy: weight_conversions.extend(get_checkpoint_conversion_mapping("legacy")) # Let the quantizer rewrite / augment the conversion pipeline. This is where the # FP8 dequantizer (when `dequantize=True`) prepends a `Fp8Dequantize` op to # every existing converter so that per-block scales are applied *before* any # expert-merge / concat ops flatten the per-expert structure away. if hf_quantizer is not None: weight_conversions = hf_quantizer.update_weight_conversions(weight_conversions) return weight_conversions