Fix non-idempotent revert_weight_conversion corrupting trust_remote_code saves#46651
Closed
Bluear7878 wants to merge 3 commits into
Closed
Fix non-idempotent revert_weight_conversion corrupting trust_remote_code saves#46651Bluear7878 wants to merge 3 commits into
Bluear7878 wants to merge 3 commits into
Conversation
…ode saves
When a model is saved without having been loaded through the conversion engine
(`model._weight_conversions is None`) -- e.g. a `trust_remote_code` model whose
remote `from_pretrained` builds the state dict directly, or a `from_config` model
-- `revert_weight_conversion` fabricates the mapping from the hardcoded table and
applies its reverse. For a mapping whose reverse rename is not idempotent (an
unanchored rename whose target is a substring of its source, e.g. nomic_bert and
jina_embeddings_v3 `WeightRenaming("encoder.layers", "layers")`), reverting keys
that are already in on-disk layout double-prepends the prefix
(`encoder.layers` -> `encoder.encoder.layers`), silently dropping every affected
weight on reload.
Only revert the fabricated mapping when the in-memory keys are actually in the
canonical (post-load) layout: if applying the forward renamings would still rename
a key, the keys are already in disk layout and must not be reverted. This preserves
the documented `from_config` -> original-format round-trip while fixing the
corruption for the whole substring-doubling family (nomic_bert, jina_embeddings_v3,
llava and aliases, emu3, colqwen2, detr).
… heuristic The previous "would the forward mapping rename any key?" check false-positived on native `from_config` models whose canonical keys still match an unanchored forward rename (e.g. dinov3_vit `(?<!model.)layer.`, qianfan_ocr `encoder.layers.`), wrongly skipping their legitimate revert and breaking `test_reverse_loading_mapping` for ~10 models. Disk-layout in-memory keys only occur for `trust_remote_code` models (the remote modeling code builds the state dict directly and never runs the conversion engine); native models always carry canonical keys. Skip the fabricated revert only when the model class lives in the dynamic `transformers_modules` namespace, leaving native models untouched.
…ration idempotent Match `transformers_modules.` with a trailing dot so an unrelated top-level module that merely starts with `transformers_modules` cannot match, and pass `overwrite=True` when registering the test conversion mappings so the tests can be re-run in the same process.
Contributor
|
CI Dashboard: View test results in Grafana |
Contributor
Author
|
@Cyrilvallez this fixes a silent |
Member
|
Thanks for the PR! The proper fix is in #46876 though! Let me know if something is unclear |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Fixes #46650
save_pretrainedsilently corrupts the checkpoint of atrust_remote_codemodel whosemodel_typehas ahardcoded conversion mapping in
conversion_mapping.py(currentlynomic_bertandjina_embeddings_v3).A remote model is loaded by its own remote modeling code, which builds the state dict directly and never
runs the core conversion engine, so
model._weight_conversions is Noneand its in-memory keys are alreadyin the original on-disk layout. On save,
revert_weight_conversion(core_model_loading.py) fabricatesthe mapping from the hardcoded table and applies its reverse anyway. For an unanchored
WeightRenamingwhose target is a substring of its source — e.g.
nomic_bert/jina_embeddings_v3WeightRenaming("encoder.layers", "layers")— reverting keys that already carry the prefix double-prependsit:
End-to-end,
AutoModel.from_pretrained("nomic-ai/nomic-embed-text-v1", trust_remote_code=True)→
save_pretrained(d)→from_pretrained(d, trust_remote_code=True)reports everyencoder.layers.*weightas missing and silently keeps the randomly-initialised weights — accuracy collapses, with no error raised.
Fix
The fabricated mapping only ever applies to a model whose in-memory keys are in the canonical (native)
layout. That is never the case for a
trust_remote_codemodel — its keys are already on-disk. So skip thefabricated revert when the model class lives in the dynamic
transformers_modulesnamespace, and revertnative models (built via
from_config, whose keys are canonical) exactly as before.Alternatives considered and ruled out:
on native models whose canonical keys still match an unanchored forward rename — it wrongly skipped their
legitimate revert and broke
test_reverse_loading_mappingfor ~10 models. The canonical-vs-disk questionis undecidable per-key for the
X.model -> Xmapping family.^encoder\.layers): only patches one model and breaks thenomic_bertwrapper classes (
base_model_prefixnesting undernomic_bert.) on both load and save.PrefixChange(à la Move some conversion mappings to PrefixChange #45567): does not address the unanchored reverse renameand still requires per-model migration.
Tests
tests/utils/test_core_model_loading.py::TestRevertWeightConversionRemoteCode:trust_remote_codemodel (class intransformers_modules) with on-disk-layout keys is saved unchanged(no
encoder.encoder.*),from_config) with canonical keys still round-trips to the on-disk layout.The full
test_reverse_loading_mappingsuite (which builds native models) is unchanged by this fix.ruff check+ruff formatclean.Code Agent Policy
Before submitting
save_pretrainedcorrupts keys (encoder.encoder.*) for trust_remote_code / from_config models via non-idempotent reverse rename #46650TestRevertWeightConversionRemoteCodeWho can review?
@Cyrilvallez — this directly extends your #45567 (
PrefixChangefilter inrevert_weight_conversion).