Skip to content

fix(checkpoint): preserve class-based mappings through model wrappers - #3749

Draft
jQizhang wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
jQizhang:larkz/fix/qwen-vl-class-mapping
Draft

jQizhang wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
jQizhang:larkz/fix/qwen-vl-class-mapping

Conversation

@jQizhang

Copy link
Copy Markdown
Contributor

Summary

Fix checkpoint key mapping for models whose Transformers conversion mappings are registered by model class rather than by config.model_type.

This primarily fixes Qwen2-VL and Qwen2.5-VL checkpoint loading, while addressing the underlying wrapper-related lookup issue for other class-registered models as well.

Problem

AutoModel's DCP-based Hugging Face checkpoint path uses get_combined_key_mapping() to build a flat regex mapping before reading safetensors metadata.

Previously, this function queried Transformers only with:

get_checkpoint_conversion_mapping(model_type)

However, Transformers registers some conversion mappings by model class name. For example:

  • Qwen2VLForConditionalGeneration
  • Qwen2_5_VLForConditionalGeneration through the Qwen2-VL class mapping alias

There is no corresponding direct mapping registered under qwen2_vl or qwen2_5_vl.

As a result, AutoModel did not apply the required Qwen-VL renames:

visual.*  -> model.visual.*
model.*   -> model.language_model.*

This becomes more subtle after FSDP2 parallelization because the runtime model may be represented by a dynamically created subclass. Looking up only type(model).__name__ would find the wrapper class rather than the original Hugging Face model class.

Transformers upstream behavior

Transformers introduced fine-grained class-based conversion mappings in PR #45661, merged as commit f397b9e.

The upstream lookup explicitly checks the concrete model class before falling back to config.model_type:

class_name = type(model).__name__
model_type = model.config.model_type

conversions = get_checkpoint_conversion_mapping(class_name)
if conversions is None and model_type:
    conversions = get_checkpoint_conversion_mapping(model_type)

That upstream change also moved model-specific mappings such as Qwen-VL to class-based registrations so that different model heads sharing the same model_type can use different checkpoint conversions.

AutoModel already consumes Transformers' conversion registry, but its simplified DCP loading path bypassed this model-aware lookup and queried only the low-level model_type mapping.

Root cause

The full Transformers conversion pipeline returns scoped WeightRenaming and WeightConverter objects and applies them while loading model weights.

AutoModel cannot directly use that pipeline in its standard DCP path because DCP metadata planning requires a flat {regex: replacement} mapping before reading tensors. AutoModel therefore extracts simple renaming rules from the Transformers registry.

During that extraction, AutoModel previously looked up only config.model_type. This diverged from the class-first behavior introduced by Transformers PR #45661 and made class-registered mappings invisible.

Calling the upstream helper with the wrapped model is not sufficient either: FSDP2 may create a dynamic subclass whose runtime class name is not registered in Transformers.

Solution

  • Pass the runtime model into get_combined_key_mapping().
  • Search the runtime class hierarchy in order, allowing transparent wrappers to resolve mappings registered for their underlying Hugging Face model class.
  • Fall back to config.model_type when no class mapping is registered.
  • Continue sourcing the actual conversion rules from Transformers rather than adding Qwen-specific mappings to AutoModel.
  • Preserve the existing precedence of AutoModel model-specific and VLM mappings.
  • Keep the new model argument optional so existing callers retain the previous behavior.

The resulting lookup order is:

runtime wrapper class
    -> underlying model classes in MRO order
    -> config.model_type

Why this is not Qwen-specific

Transformers also registers mappings by class for models such as:

  • LlavaModel
  • ViTModel
  • CLIPVisionModel
  • Qwen2AudioModel
  • DetrModel

Any of these mappings can become invisible when a transparent framework wrapper changes the runtime class name.

Resolving the underlying class fixes the general wrapper compatibility issue without maintaining a separate model allowlist or duplicating Transformers conversion rules in AutoModel.

Testing

Added coverage for:

  • Qwen2.5-VL class-registered mappings.
  • Qwen2-VL class-registered mappings.
  • Dynamically created wrapper subclasses.
  • Vision and language tower key renaming.
  • Avoiding repeated model.language_model prefixes.
  • End-to-end loading from a safetensors checkpoint into the expected in-memory model parameters.

Validation:

Before fix: 2 failed, 1 passed
After fix:  4 passed

The following checks also pass:

ruff check:        passed
ruff format check: passed

Signed-off-by: larkzhang-nv <larkz@nvidia.com>
@copy-pr-bot

copy-pr-bot Bot commented Aug 29, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant