Fix Transformers 4.x/5.x compatibility and isolate multi-image vision attention

#102
by Bvicii - opened

Makes the remote code load, generate and save with trust_remote_code=True on Transformers 4.54 through 5.22, without a version pin.

Changes

  • Causal mask: always use transformers.masking_utils.create_causal_mask. Its signature differences across releases (input_embeds → inputs_embeds, removal of cache_position) are resolved once at import. This keeps the SDPA is_causal fast path, the flex-attention BlockMask, and the FA2/FA3/FA4 mask handling. Only dense 4D tensor masks are converted locally; a BlockMask passed in by the caller goes to the native function, which returns it unchanged.
  • Default RoPE: v5 no longer has default in ROPE_INIT_FUNCTIONS, so the rotary embedding computes default parameters itself.
  • Initialization and derived buffers: v5 builds the model on the meta device. Weights that are already loaded are no longer re-initialized, and non-persistent buffers (rotary inv_freq / original_inv_freq, vision position_ids) are recomputed after loading.
  • Flash Attention 2 on v5: declare _supports_flash_attn = True. v5 only reads this name, so attn_implementation="flash_attention_2" used to fail with "PaddleOCRVisionModel does not support Flash Attention 2 yet".
  • save_pretrained on v5: _tied_weights_keys is now a dict ({"lm_head.weight": "model.embed_tokens.weight"}). The list form raised AttributeError: 'list' object has no attribute 'keys'. tie_word_embeddings stays false, so lm_head is still untied.
  • Generation: prepare_inputs_for_generation builds cache_position itself when v5 doesn't pass it in. It now drops pixel_values only after the KV cache holds the prompt, so use_cache=False generation keeps the image features.

Behavior change: multi-image inputs

When several images are packed into one request, eager and SDPA vision attention now run per image, using cu_seqlens. Before this change, patches of one image could attend to patches of another image, so outputs for multi-image requests with eager/SDPA will change. They now match encoding each image on its own, which is what the FA2 path (flash_attn_varlen_func) already did. Single-image outputs and FA2 outputs are unchanged.

A vision padding mask is cut into one block per image. Broadcastable masks such as [B, 1, 1, L] keep their size-1 dims, so only the real sequence dims are sliced.

Testing

  • The earlier commits were validated with the full checkpoint on sampled Transformers 4.x/5.x releases (CPU float32).
  • The latest commits were tested on a tiny random-weight model with Transformers 4.57.6, 5.12.1 and 5.22.1, on CPU float32 (eager, SDPA) and A100 bfloat16 (eager, SDPA, FA2). The checks were:
    • generation matches the previous PR head
    • generation with cache matches generation without cache
    • a left-padded row matches the same row run without padding
    • packed multi-image vision encoding matches encoding each image separately
    • packed vision attention with [B,1,1,L] and [B,1,L,L] padding masks matches running each image alone (eager, SDPA)
    • a flex-attention BlockMask comes back from the text mask path as the same object
    • save_pretrained + from_pretrained restores the buffers, keeps lm_head tying as configured, and reproduces pre-save generation
  • Numerical alignment with the native Transformers implementation (trust_remote_code=False), using the PaddleOCR-VL-1.6 checkpoint on A100 bfloat16 with a two-image OCR request:
    • On Transformers 5.12.1, vision features from this code are bit-identical to the native implementation with eager and SDPA. FA2 differs by about 6% relative L2, because this code keeps its own flash_attn_varlen_func path. That is about the same size as the eager vs SDPA difference within one implementation.
    • On 5.22.1, native vision features differ from 5.12.1 by about 16%. Since 5.13.0, the native position-embedding interpolation is equivalent to align_corners=True (5.15.0+ sets interpolation_align_corners = True), while the original model and this code use align_corners=False. With the native flag set to False, native and this code match within 5e-6 relative L2 (float32 eager). This code produces bit-identical features on 5.12.1 and 5.22.1.
    • The OCR text is identical across native and this code, all three backends, and both versions.
    • Packed two-image features equal per-image features, which matches the native implementation. With the old joint attention, eager/SDPA mixed text from the second image into the first image's OCR output.

The same change is applied to PaddleOCR-VL (#102), PaddleOCR-VL-1.5 (#35) and PaddleOCR-VL-1.6 (#15).

🤖 Generated with Claude Code

Bvicii changed pull request title from Fix remote-code compatibility across Transformers 4.x and 5.x to Fix Transformers 4.x/5.x compatibility and isolate multi-image vision attention
Publish this branch
This branch is in draft mode, publish it to be able to merge.

Sign up or log in to comment