fix(loading/paddleocr-vl): remap real language_model.* / visual.* checkpoint keys - #607
Merged
Conversation
The published PaddleOCR-VL checkpoint (PaddleOCRVLForConditionalGeneration) wraps the ERNIE-4.5 text backbone under `language_model` (`language_model.model.*`, `language_model.lm_head.*`), so the loader's remap never produced the `model.*` / `lm_head.*` keys the text model requests. Loading the real checkpoint failed with `Failed to load PaddleOCR-VL text model: Weight not found: model.embed_tokens.weight`. Strip the `language_model.` wrapper in `remap_key` so the ERNIE-4.5 backbone sees `model.*` / `lm_head.*`. The published vision tower is already sanitized and qkv-fused under `visual.*` (embeddings, layers, post_layernorm, projector), which the encoder and connector already consume as pass-through, so no vision-side change is needed. The existing reference `Model.sanitize` handling (`visual.vision_model.*` to `visual.*`, split `q_proj`/`k_proj`/`v_proj` to fused `qkv`, `mlp_AR` to `visual.projector`) is preserved for split exports. Add `vlm_paddleocr_tests.rs`: feed the real checkpoint key set (read read-only from `model.safetensors.index.json`) through the remap and assert the produced keys match exactly what the text model, vision encoder, and connector request, that no `language_model.` wrapper survives, and that the remap is lossless. Also cover `remap_key` for the wrapper strip, sanitized-vision pass-through, and the legacy reference layout.
16 tasks
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
1 participant
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.
Follow-up hardening for #526: the merged PaddleOCR-VL port cannot load its real checkpoint.
mlxcel generate -m /home/inureyes/models/paddleocr-vl-bfloat16 ...fails withFailed to load PaddleOCR-VL text model: Weight not found: model.embed_tokens.weight.Root cause
The published checkpoint (
PaddleOCRVLForConditionalGeneration) has exactly two top-level weight prefixes,language_model.*andvisual.*:language_model.model.embed_tokens.weight,language_model.model.layers.{i}.*,language_model.model.norm.weight,language_model.lm_head.weightvisual.embeddings.*,visual.layers.{i}.*(attention already fused asself_attn.qkv),visual.post_layernorm.*,visual.projector.*The loader's
remap_keywas written for the referenceModel.sanitizelayout (visual.vision_model.*, splitq_proj/k_proj/v_proj, top-levelmlp_AR, baremodel.*/lm_head.*). None of those patterns appear in the real checkpoint (verified read-only: zero keys match them), so the text keys passed through unstripped and the ERNIE-4.5 backbone could not findmodel.embed_tokens.weight.Fix
Strip the
language_model.wrapper inremap_keysolanguage_model.model.*becomesmodel.*andlanguage_model.lm_head.*becomeslm_head.*. The vision tower is already sanitized and qkv-fused undervisual.*, matching exactly whatPaddleOcrVisionEncoder::from_weights(prefix = "visual")andPaddleOcrProjector::from_weights(prefix = "visual.projector")consume, so it passes through unchanged. The reference-layout branches (visual.vision_model.*tovisual.*, split-qkv fusion,mlp_ARtovisual.projector) are preserved for split exports.What changed
src/loading/vlm_paddleocr.rs:remap_keystrips thelanguage_model.prefix; theremap_paddleocr_weightsdoc comment now describes both the published and reference layouts.src/loading/vlm_paddleocr_tests.rs(new): feeds the real checkpoint key set (18 text layers, 27 vision layers, as read read-only frommodel.safetensors.index.json) through the remap and asserts the produced keys match exactly what the text model, vision encoder, and connector request; asserts nolanguage_model.wrapper survives and that the remap is lossless; plus focusedremap_keycases for the wrapper strip, sanitized-vision pass-through, and the legacy reference layout.Test plan
cargo test --lib paddleocr(5 passed)cargo clippy --lib --tests -- -D warnings(clean)cargo check --lib --tests(clean)cargo fmt --check(clean)Checkpoint keys were inspected read-only from
model.safetensors.index.json; the on-GPU real-model load is re-validated by the orchestrator.