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

modeling_paddleocr_vl.py

@@ -12,9 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import inspect
 import math
 import warnings
 from dataclasses import dataclass
+from functools import wraps
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import numpy as np
@@ -28,9 +30,9 @@ from transformers.cache_utils import (
     Cache,
     DynamicCache,
 )
+from transformers.masking_utils import create_causal_mask as _hf_create_causal_mask
 from transformers.generation import GenerationMixin
 from transformers.integrations import use_kernel_forward_from_hub
-from transformers.masking_utils import create_causal_mask
 from transformers.modeling_attn_mask_utils import AttentionMaskConverter
 from transformers.modeling_layers import GradientCheckpointingLayer
 from transformers.modeling_outputs import (
@@ -66,6 +68,87 @@ else:
 from .configuration_paddleocr_vl import PaddleOCRVisionConfig, PaddleOCRVLConfig
 
 
+# Resolve optional framework interfaces once, outside the forward path.
+try:
+    from transformers import initialization as _hf_initialization
+except ImportError:
+    _hf_initialization = None
+
+# create_causal_mask exists in every release this file can import on (>= 4.54),
+# but its signature changed: `input_embeds` -> `inputs_embeds`, and
+# `cache_position` was later removed in favor of reading offsets from the cache.
+_HF_MASK_PARAMETERS = frozenset(inspect.signature(_hf_create_causal_mask).parameters)
+_HF_MASK_INPUT_NAME = "inputs_embeds" if "inputs_embeds" in _HF_MASK_PARAMETERS else "input_embeds"
+
+
+def _guard_initializer(function):
+    @wraps(function)
+    def initialize(tensor, *args, **kwargs):
+        if getattr(tensor, "_is_hf_initialized", False):
+            return tensor
+        with torch.no_grad():
+            return function(tensor, *args, **kwargs)
+    return initialize
+
+
+@_guard_initializer
+def _init_tensor(tensor, operation, *args, **kwargs):
+    # Standard primitives use the framework's protected implementation when
+    # present. Keep model-specific TF/LeCun sampling formulas below unchanged.
+    function = getattr(_hf_initialization, operation, None)
+    if function is not None:
+        return function(tensor, *args, **kwargs)
+    if operation == "copy_":
+        return tensor.copy_(*args, **kwargs)
+    return getattr(nn.init, operation)(tensor, *args, **kwargs)
+
+
+def _create_causal_mask(config, inputs_embeds, attention_mask, past_key_values, cache_position):
+    # Only dense tensors need conversion; flex BlockMask is passed through to the native function.
+    if isinstance(attention_mask, torch.Tensor) and attention_mask.ndim == 4:
+        mask = attention_mask.to(device=inputs_embeds.device)
+        if mask.dtype == torch.bool:
+            additive = torch.zeros(mask.shape, device=mask.device, dtype=inputs_embeds.dtype)
+            return additive.masked_fill(~mask, torch.finfo(inputs_embeds.dtype).min)
+        return mask.to(dtype=inputs_embeds.dtype)
+
+    mask_kwargs = {
+        "config": config,
+        _HF_MASK_INPUT_NAME: inputs_embeds,
+        "attention_mask": attention_mask,
+        "past_key_values": past_key_values,
+    }
+    if "cache_position" in _HF_MASK_PARAMETERS:
+        mask_kwargs["cache_position"] = cache_position
+    if "position_ids" in _HF_MASK_PARAMETERS:
+        # 3D M-RoPE coordinates are not packed text-sequence positions.
+        mask_kwargs["position_ids"] = None
+    return _hf_create_causal_mask(**mask_kwargs)
+
+
+def _init_derived_buffers(module):
+    # v5 materializes non-persistent buffers from meta storage; they are not
+    # loaded from the checkpoint and must be initialized explicitly.
+    if isinstance(module, RotaryEmbedding):
+        inv_freq, module.attention_scaling = module.rope_init_fn(
+            module.config, device=module.inv_freq.device
+        )
+        _init_tensor(module.inv_freq, "copy_", inv_freq)
+        module.original_inv_freq = module.inv_freq
+    elif isinstance(module, SigLIPRotaryEmbedding):
+        inv_freq = 1.0 / (
+            module.theta ** (
+                torch.arange(0, module.dim, 2, dtype=torch.float32, device=module.inv_freq.device)
+                / module.dim
+            )
+        )
+        _init_tensor(module.inv_freq, "copy_", inv_freq)
+    elif isinstance(module, PaddleOCRVisionEmbeddings):
+        _init_tensor(module.position_ids, "copy_",
+            torch.arange(module.num_positions, device=module.position_ids.device).expand(1, -1)
+        )
+
+
 class RotaryEmbedding(nn.Module):
     def __init__(self, config: PaddleOCRVLConfig, device=None):
         super().__init__()
@@ -92,12 +175,25 @@ class RotaryEmbedding(nn.Module):
         self.original_max_seq_len = config.max_position_embeddings
 
         self.config = config
-        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+        if self.rope_type == "default":
+            self.rope_init_fn = self.compute_default_rope_parameters
+        else:
+            self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
 
         inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
         self.register_buffer("inv_freq", inv_freq, persistent=False)
         self.original_inv_freq = self.inv_freq
 
+    @staticmethod
+    def compute_default_rope_parameters(config, device=None, seq_len=None):
+        # Default RoPE is no longer part of ROPE_INIT_FUNCTIONS in v5.
+        head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+        dim = int(head_dim * getattr(config, "partial_rotary_factor", 1.0))
+        inv_freq = 1.0 / (
+            config.rope_theta ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim)
+        )
+        return inv_freq, 1.0
+
     def _dynamic_frequency_update(self, position_ids, device):
         """
         dynamic RoPE layers should recompute `inv_freq` in the following situations:
@@ -539,6 +635,12 @@ class Ernie4_5PreTrainedModel(PreTrainedModel):
     }
 
 
+    @torch.no_grad()
+    def _init_weights(self, module):
+        super()._init_weights(module)
+        _init_derived_buffers(module)
+
+
 @auto_docstring
 class Ernie4_5Model(Ernie4_5PreTrainedModel):
     def __init__(self, config: PaddleOCRVLConfig):
@@ -603,13 +705,8 @@ class Ernie4_5Model(Ernie4_5PreTrainedModel):
         elif position_ids.dim() == 2:
             position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
 
-        causal_mask = create_causal_mask(
-            config=self.config,
-            inputs_embeds=inputs_embeds,
-            attention_mask=attention_mask,
-            past_key_values=past_key_values,
-            position_ids=position_ids,
-            cache_position=cache_position,
+        causal_mask = _create_causal_mask(
+            self.config, inputs_embeds, attention_mask, past_key_values, cache_position
         )
 
         hidden_states = inputs_embeds
@@ -634,7 +731,7 @@ class Ernie4_5Model(Ernie4_5PreTrainedModel):
 
 
 class Ernie4_5ForCausalLM(Ernie4_5PreTrainedModel, GenerationMixin):
-    _tied_weights_keys = ["lm_head.weight"]
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
     _tp_plan = {"lm_head": "colwise_rep"}
     _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
 
@@ -711,6 +808,7 @@ class Ernie4_5ForCausalLM(Ernie4_5PreTrainedModel, GenerationMixin):
         )
 
 
+@_guard_initializer
 def _trunc_normal_(tensor, mean, std, a, b):
     # Cut & paste from PyTorch official master until it's in a few official releases - RW
     # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
@@ -747,6 +845,7 @@ def _trunc_normal_(tensor, mean, std, a, b):
     tensor.clamp_(min=a, max=b)
 
 
+@_guard_initializer
 def trunc_normal_tf_(
     tensor: torch.Tensor,
     mean: float = 0.0,
@@ -777,6 +876,7 @@ def trunc_normal_tf_(
         tensor.mul_(std).add_(mean)
 
 
+@_guard_initializer
 def variance_scaling_(tensor, scale=1.0, mode="fan_in", distribution="normal"):
     fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
     if mode == "fan_in":
@@ -1134,16 +1234,39 @@ class PaddleOCRAttention(nn.Module):
                 elif self.config._attn_implementation == "sdpa":
                     attention_interface = sdpa_attention_forward
 
-            attn_output, attn_weights = attention_interface(
-                self,
-                queries,
-                keys,
-                values,
-                attention_mask,
-                is_causal=self.is_causal,
-                scaling=self.scale,
-                dropout=0.0 if not self.training else self.dropout,
-            )
+            if cu_seqlens is not None and cu_seqlens.numel() > 2:
+                boundaries = cu_seqlens.tolist()
+                outputs = []
+                attn_weights = (
+                    queries.new_zeros(batch_size, self.num_heads, seq_length, seq_length)
+                    if output_attentions else None
+                )
+                for start, end in zip(boundaries[:-1], boundaries[1:]):
+                    local_mask = None
+                    if attention_mask is not None:
+                        # Keep size-1 broadcast dims, e.g. a [B, 1, 1, L] padding mask.
+                        span = slice(start, end)
+                        local_mask = attention_mask[
+                            ...,
+                            span if attention_mask.shape[-2] > 1 else slice(None),
+                            span if attention_mask.shape[-1] > 1 else slice(None),
+                        ]
+                    output, weights = attention_interface(
+                        self, queries[:, :, start:end], keys[:, :, start:end],
+                        values[:, :, start:end], local_mask,
+                        is_causal=False, scaling=self.scale,
+                        dropout=0.0 if not self.training else self.dropout,
+                    )
+                    outputs.append(output)
+                    if attn_weights is not None and weights is not None:
+                        attn_weights[..., start:end, start:end] = weights
+                attn_output = torch.cat(outputs, dim=1)
+            else:
+                attn_output, attn_weights = attention_interface(
+                    self, queries, keys, values, attention_mask,
+                    is_causal=self.is_causal, scaling=self.scale,
+                    dropout=0.0 if not self.training else self.dropout,
+                )
             attn_output = attn_output.reshape(
                 batch_size, seq_length, embed_dim
             ).contiguous()
@@ -1263,45 +1386,47 @@ class PaddleOCRPreTrainedModel(PreTrainedModel):
         "PaddleOCRVisionEmbeddings",
         "PaddleOCRMultiheadAttentionPoolingHead",
     ]
-    _supports_flash_attn_2 = True
+    # v5 only reads `_supports_flash_attn`; 4.x accepts both names.
+    _supports_flash_attn = True
     _supports_sdpa = True
 
     def _init_weights(self, module):
-        """Initialize the weights"""
+        """Initialize the weights and derived non-persistent buffers."""
+        _init_derived_buffers(module)
         if isinstance(module, PaddleOCRVisionEmbeddings):
             width = (
                 self.config.vision_config.hidden_size
                 if isinstance(self.config, PaddleOCRVLConfig)
                 else self.config.hidden_size
             )
-            nn.init.normal_(module.position_embedding.weight, std=1 / np.sqrt(width))
+            _init_tensor(module.position_embedding.weight, 'normal_', std=1 / np.sqrt(width))
         elif isinstance(module, nn.Embedding):
             default_flax_embed_init(module.weight)
         elif isinstance(module, PaddleOCRAttention):
-            nn.init.xavier_uniform_(module.q_proj.weight)
-            nn.init.xavier_uniform_(module.k_proj.weight)
-            nn.init.xavier_uniform_(module.v_proj.weight)
-            nn.init.xavier_uniform_(module.out_proj.weight)
-            nn.init.zeros_(module.q_proj.bias)
-            nn.init.zeros_(module.k_proj.bias)
-            nn.init.zeros_(module.v_proj.bias)
-            nn.init.zeros_(module.out_proj.bias)
+            _init_tensor(module.q_proj.weight, 'xavier_uniform_')
+            _init_tensor(module.k_proj.weight, 'xavier_uniform_')
+            _init_tensor(module.v_proj.weight, 'xavier_uniform_')
+            _init_tensor(module.out_proj.weight, 'xavier_uniform_')
+            _init_tensor(module.q_proj.bias, 'zeros_')
+            _init_tensor(module.k_proj.bias, 'zeros_')
+            _init_tensor(module.v_proj.bias, 'zeros_')
+            _init_tensor(module.out_proj.bias, 'zeros_')
         elif isinstance(module, PaddleOCRMLP):
-            nn.init.xavier_uniform_(module.fc1.weight)
-            nn.init.xavier_uniform_(module.fc2.weight)
-            nn.init.normal_(module.fc1.bias, std=1e-6)
-            nn.init.normal_(module.fc2.bias, std=1e-6)
+            _init_tensor(module.fc1.weight, 'xavier_uniform_')
+            _init_tensor(module.fc2.weight, 'xavier_uniform_')
+            _init_tensor(module.fc1.bias, 'normal_', std=1e-6)
+            _init_tensor(module.fc2.bias, 'normal_', std=1e-6)
         elif isinstance(module, PaddleOCRMultiheadAttentionPoolingHead):
-            nn.init.xavier_uniform_(module.probe.data)
-            nn.init.xavier_uniform_(module.attention.in_proj_weight.data)
-            nn.init.zeros_(module.attention.in_proj_bias.data)
+            _init_tensor(module.probe, 'xavier_uniform_')
+            _init_tensor(module.attention.in_proj_weight, 'xavier_uniform_')
+            _init_tensor(module.attention.in_proj_bias, 'zeros_')
         elif isinstance(module, (nn.Linear, nn.Conv2d)):
             lecun_normal_(module.weight)
             if module.bias is not None:
-                nn.init.zeros_(module.bias)
+                _init_tensor(module.bias, 'zeros_')
         elif isinstance(module, nn.LayerNorm):
-            module.bias.data.zero_()
-            module.weight.data.fill_(1.0)
+            _init_tensor(module.bias, 'zeros_')
+            _init_tensor(module.weight, 'ones_')
 
 
 # Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoder with AltCLIP->PaddleOCR
@@ -1866,7 +1991,7 @@ class PaddleOCRVLCausalLMOutputWithPast(ModelOutput):
 
 
 class PaddleOCRVLForConditionalGeneration(Ernie4_5PreTrainedModel, GenerationMixin):
-    _tied_weights_keys = ["lm_head.weight"]
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
     config_class = PaddleOCRVLConfig
     _no_split_modules = ["Ernie4_5_DecoderLayer", "PaddleOCREncoderLayer"]
 
@@ -2342,7 +2467,21 @@ class PaddleOCRVLForConditionalGeneration(Ernie4_5PreTrainedModel, GenerationMix
 
         model_inputs["position_ids"] = None
 
-        if cache_position[0] != 0:
+        # v5 creates cache_position inside the parent method, rather than
+        # passing it to this override. Read the prepared value back first.
+        cache_position = model_inputs.get("cache_position", cache_position)
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            prepared = model_inputs.get("inputs_embeds")
+            if prepared is None:
+                prepared = model_inputs["input_ids"]
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + prepared.shape[1], device=prepared.device
+            )
+            model_inputs["cache_position"] = cache_position
+
+        # Without KV caching, every forward still consumes the image tokens.
+        if use_cache and past_key_values is not None and past_key_values.get_seq_length() > 0:
             model_inputs["pixel_values"] = None
             model_inputs["pixel_values_videos"] = None