Repository navigation
Improve nested base_model_prefix handling in weight conversion and loading - #45421
yonigozlan wants to merge 5 commits into
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
|
@Cyrilvallez does #45448 also handle this case? I think we now have 2 cases with the prefix
Unsure about the current state and who is doing what |
|
[For maintainers] Suggested jobs to run (before merge) run-slow: conditional_detr, detr |
| def _strip_model_prefix_for_save(key: str, model) -> str: | ||
| """ | ||
| Recursively strip all `base_model_prefix` segments from a state-dict key so that | ||
| reverse conversion rules (written relative to the innermost base model) operate on | ||
| bare keys regardless of nesting depth. | ||
|
|
||
| We identify each prefix level by finding the direct child whose `base_model_prefix` | ||
| matches its attribute name (same logic as `_compute_all_prefixes`). | ||
|
|
||
| Examples for `DetrForSegmentation` (prefix chain `detr` -> `model`): | ||
|
|
||
| "detr.model.backbone.x" -> "backbone.x" | ||
| "detr.class_labels_classifier.x" -> "class_labels_classifier.x" | ||
| "mask_head.x" -> "mask_head.x" | ||
| """ | ||
| for name, child in model.named_children(): | ||
| child_prefix = getattr(child, "base_model_prefix", "") | ||
| if child_prefix and child_prefix == name and key.startswith(name + "."): | ||
| stripped_key = key[len(name) + 1 :] | ||
| return _strip_model_prefix_for_save(stripped_key, child) | ||
| return key |
There was a problem hiding this comment.
This might be made redundant by #45448
The rest I think is somewhat orthogonal, could you confirm @Cyrilvallez ?
| # Opt in via `_checkpoint_conversion_prefix_free = True` when the source checkpoint is fully flat, | ||
| # so that all prefixes should be stripped before saving. | ||
| strip_prefix = getattr(model, "_checkpoint_conversion_prefix_free", False) | ||
|
|
||
| state_dict = sorted(state_dict.items(), key=lambda kv: dot_natural_key(kv[0])) | ||
| for original_key, tensor in state_dict: | ||
| # Rename the key according to all renaming pattern and optional weight converter patterns | ||
| renamed_key, source_pattern = rename_source_key(original_key, renamings, converters) | ||
| bare_key = _strip_model_prefix_for_save(original_key, model) if strip_prefix else original_key | ||
| renamed_key, source_pattern = rename_source_key(bare_key, renamings, converters) | ||
| if source_pattern is not None: | ||
| new_converter = deepcopy(pattern_to_converter[source_pattern]) | ||
| # each target key gets its own converter instance | ||
| mapping = conversion_mapping.setdefault(renamed_key, new_converter) | ||
| else: | ||
| mapping = conversion_mapping.setdefault(renamed_key, WeightRenaming(original_key, renamed_key)) | ||
| source_pattern = original_key | ||
| mapping = conversion_mapping.setdefault(renamed_key, WeightRenaming(bare_key, renamed_key)) | ||
| source_pattern = bare_key |
Cyrilvallez
left a comment
Cyrilvallez
left a comment
There was a problem hiding this comment.
Hey @yonigozlan! Concerning this
Loading with more than one prefix level: rename_source_key step 3 only tried a single model.base_model_prefix to add or remove. Loading a doubly-nested model correctly was not possible without hacks. That means for example we couldn't load weights from DetrModel in DetrModelForSegmentation, as we need to support adding nested prefixes for this to work.
The base_model_prefix is not intended to be used for conversions! It's only there to allow loading a BaseModel from a ForXXX checkpoint, and the opposite, and find the embeddings correctly etc.
We should not use them to replace the conversion mappings.
Indeed, you should be able to use PrefixChange from #45448 to help you here I believe! Let me know if you need help/you encounter issues!
| base_model_prefix = "conditional_detr" | ||
|
|
There was a problem hiding this comment.
We cannot have that, the base model is named model
There was a problem hiding this comment.
ConditionalDetrForObjectDetection is used as a "base model" in ConditionalDetrForSegmentation as self.detr, hence why we need to modify it here. The modifications in core_model_loading.py now fully works with loading/saving these kind of nested models
| base_model_prefix = "detr" | ||
|
|
There was a problem hiding this comment.
Same here, DetrForObjectDetection is used as a "base_model" in DetrForSegmentation
| # 3. check if we need to add or remove prefix if necessary (only during loading, not saving) | ||
| if prefix is not None and meta_state_dict is not None: | ||
| if ( | ||
| renamed_key.startswith(prefix) | ||
| and meta_state_dict.get(re.sub(f"^{prefix}.", "", renamed_key, count=1)) is not None | ||
| ): | ||
| renamed_key = re.sub(f"^{prefix}.", "", renamed_key, count=1) | ||
| elif meta_state_dict.get(f"{prefix}.{renamed_key}") is not None: | ||
| renamed_key = f"{prefix}.{renamed_key}" | ||
| # 3. If the key is still not in the model state dict, try adding or removing each | ||
| # prefix level (longest first) until a match is found. Only active during loading. | ||
| if valid_prefixes is not None and meta_state_dict is not None and renamed_key not in meta_state_dict: | ||
| renamed_key = _resolve_key_for_prefix_nesting(renamed_key, valid_prefixes, meta_state_dict) |
There was a problem hiding this comment.
This is only used to be able to load BaseModel from ForXXXModel and opposite. It's really not intended for anything else, this should be only the conversion mappings! Very important to not confuse both, as it will otherwise become very quickly unmaintanable
There was a problem hiding this comment.
Yes agreed! But we still need this to be able to support nested ForXXXModel, and load each level of weights in each level of nesting
ArthurZucker
left a comment
ArthurZucker
left a comment
There was a problem hiding this comment.
I agree with Cyril here sorry for the back and forth, I think we should NOT allow this in general its too big, too hard to maintain and you have too many cases to handle.
It will gracefully error out, let's add a ref to doc on how to load a model CLIP from a DETRForObjectDetection for example (manually pass prefix change for example) + just allow it for detr! 🤗
| pass | ||
|
|
||
|
|
||
| def _compute_all_prefixes(model) -> list[str]: |
There was a problem hiding this comment.
We don't need this AFAIK:
Upon __init__ of the parent DetrForSegmentation, it should have access to self.prefixes which can be obtain using the recursive nature of __init__. This prevent having to go over the model again, and is only run once at init time.
| valid_prefixes: list[str], | ||
| meta_state_dict: dict, | ||
| ) -> str: | ||
| """ |
There was a problem hiding this comment.
longest first can be done with regex piping (shortex | short | sh) will go with longest first we don't need all this engineering no?
|
Closing per the discussion with @Cyrilvallez and @ArthurZucker: we don't want to support multi-level |
Split off from the RF-DETR PR (#36895). These are changes to
core_model_loading.pyandtest_modeling_common.pythat are needed to load and save RF-DETR checkpoints correctly and in the same format as the original roboflow checkpoints, but they fix pre-existing issues in transformers that affect any doubly-nested model.Why RF-DETR needs this
RfDetrForInstanceSegmentationhas a two-level nesting: it wrapsRfDetrForObjectDetection(asself.model) which wrapsRfDetrModel(asself.model). The HF state dict therefore has keys undermodel.model.*. The conversion rules inconversion_mapping.pyare written relative toRfDetrModel(bare keys, no prefix), matching Roboflow's checkpoint format. So if we the saved checkpoint to be compatible with roboflow original ones, we need to be able to add nested prefixes when loading, and strip nested prefixes when saving.What was already broken
Loading with more than one prefix level:
rename_source_keystep 3 only tried a singlemodel.base_model_prefixto add or remove. Loading a doubly-nested model correctly was not possible without hacks. That means for example we couldn't load weights from DetrModel in DetrModelForSegmentation, as we need to support adding nested prefixes for this to work.Saving without prefixes: To have a more comprehensive support of source weights with the weight converter, this PR adds support for saving models without nesting, if the original checkpoints were flat.
Changes
_compute_all_prefixes(model): walks thebase_model_prefix/ submodel chain and returns all cumulative prefix paths, e.g.["model"]forLlamaForCausalLM,["model", "model.model"]forRfDetrForInstanceSegmentation.rename_source_keynow takes this list instead of a single prefix string and tries each level (longest first) to add or strip when resolving a key against the model state dict._strip_model_prefix_for_save(key, model): recursively strips all prefix levels from a key before applying reverse conversion rules. Opt-in via_checkpoint_conversion_prefix_free = Trueon the model class, so it only fires when the source checkpoint genuinely has bare keys.One other somewhat related fix:
test_reverse_loading_mapping: Now precomputes state-dict keys for all classes inall_model_classesbefore the loop. Instead of failing everytime a rule isn't used in all classes, it now asserts the rule matches at least one other class.