Update modeling_opensci.py

modeling_opensci.py

@@ -27,18 +27,19 @@ from transformers.activations import ACT2FN
 from transformers.cache_utils import Cache, DynamicCache, StaticCache
 from transformers.generation import GenerationMixin
 from transformers.modeling_attn_mask_utils import AttentionMaskConverter
-from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+# from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
 from transformers.modeling_outputs import (
     BaseModelOutputWithPast,
     CausalLMOutputWithPast,
+    QuestionAnsweringModelOutput,
     SequenceClassifierOutputWithPast,
+    TokenClassifierOutput,
 )
 from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
 from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
-from transformers.processing_utils import Unpack
 from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
 from transformers.utils import (
-    TransformersKwargs,
+    add_code_sample_docstrings,
     add_start_docstrings,
     add_start_docstrings_to_model_forward,
     logging,
@@ -47,6 +48,20 @@ from transformers.utils import (
 from transformers.utils.deprecation import deprecate_kwarg
 from .configuration_opensci import OpensciConfig
 
+# OELLM_PATCHED_v2
+import transformers.utils as _tu_oellm
+if hasattr(_tu_oellm, 'TransformersKwargs'):
+    TransformersKwargs = _tu_oellm.TransformersKwargs
+else:
+    from typing import TypedDict as _TypedDict_oellm
+    class TransformersKwargs(_TypedDict_oellm, total=False): pass
+import transformers.processing_utils as _tp_oellm
+if hasattr(_tp_oellm, 'Unpack'):
+    Unpack = _tp_oellm.Unpack
+else:
+    from typing_extensions import Unpack
+
+
 
 logger = logging.get_logger(__name__)
 
@@ -88,7 +103,20 @@ class OpensciRotaryEmbedding(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 in ROPE_INIT_FUNCTIONS:
+            self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+        else:
+            # 'default' rope type was removed from ROPE_INIT_FUNCTIONS in transformers 5.0;
+            # inline the standard computation as a fallback.
+            def _default_rope_init(config, device=None, **kwargs):
+                base = config.rope_theta
+                head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+                dim = int(head_dim * getattr(config, "partial_rotary_factor", 1.0))
+                return (
+                    1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim)),
+                    1.0,
+                )
+            self.rope_init_fn = _default_rope_init
 
         inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
         self.register_buffer("inv_freq", inv_freq, persistent=False)
@@ -118,6 +146,14 @@ class OpensciRotaryEmbedding(nn.Module):
         if "dynamic" in self.rope_type:
             self._dynamic_frequency_update(position_ids, device=x.device)
 
+        # Re-initialize inv_freq if it was not properly materialized (transformers 5.x meta-device loading).
+        # Non-persistent buffers are not stored in the checkpoint; 5.x leaves them as uninitialized
+        # memory after materialization (can be zeros, denormals, or large garbage values).
+        if not getattr(self, "_inv_freq_ready", False):
+            inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device=x.device)
+            self.register_buffer("inv_freq", inv_freq.to(x.device), persistent=False)
+            self._inv_freq_ready = True
+
         # Core RoPE block
         inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
         position_ids_expanded = position_ids[:, None, :].float()
@@ -262,7 +298,7 @@ class OpensciAttention(nn.Module):
         attention_mask: Optional[torch.Tensor],
         past_key_value: Optional[Cache] = None,
         cache_position: Optional[torch.LongTensor] = None,
-        **kwargs: Unpack[FlashAttentionKwargs],
+        # **kwargs: Unpack[FlashAttentionKwargs],
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
         input_shape = hidden_states.shape[:-1]
         hidden_shape = (*input_shape, -1, self.head_dim)
@@ -283,14 +319,24 @@ class OpensciAttention(nn.Module):
             key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
 
         attention_interface: Callable = eager_attention_forward
-        if self.config._attn_implementation != "eager":
-            if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False):
-                logger.warning_once(
-                    "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to "
-                    'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
+        if self.config._attn_implementation == "sdpa":
+            # Inline the 4.x-compatible SDPA wrapper: never passes enable_gqa so the
+            # same PyTorch kernel is selected regardless of transformers version.
+            def attention_interface(module, query, key, value, attention_mask, scaling, dropout=0.0, is_causal=None, **kw):
+                if hasattr(module, "num_key_value_groups"):
+                    key = repeat_kv(key, module.num_key_value_groups)
+                    value = repeat_kv(value, module.num_key_value_groups)
+                query, key, value = query.contiguous(), key.contiguous(), value.contiguous()
+                if is_causal is None:
+                    is_causal = attention_mask is None and query.shape[2] > 1
+                out = torch.nn.functional.scaled_dot_product_attention(
+                    query, key, value, attn_mask=attention_mask,
+                    dropout_p=dropout, scale=scaling, is_causal=is_causal,
                 )
-            else:
-                attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+                return out.transpose(1, 2).contiguous(), None
+        elif self.config._attn_implementation not in ("eager", "sdpa"):
+            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
 
         attn_output, attn_weights = attention_interface(
             self,
@@ -300,7 +346,7 @@ class OpensciAttention(nn.Module):
             attention_mask,
             dropout=0.0 if not self.training else self.attention_dropout,
             scaling=self.scaling,
-            **kwargs,
+            # **kwargs,
         )
 
         attn_output = attn_output.reshape(*input_shape, -1).contiguous()
@@ -329,7 +375,7 @@ class OpensciDecoderLayer(nn.Module):
         use_cache: Optional[bool] = False,
         cache_position: Optional[torch.LongTensor] = None,
         position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC
-        **kwargs: Unpack[FlashAttentionKwargs],
+        # **kwargs: Unpack[FlashAttentionKwargs],
     ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
         residual = hidden_states
 
@@ -339,13 +385,13 @@ class OpensciDecoderLayer(nn.Module):
         hidden_states, self_attn_weights = self.self_attn(
             hidden_states=hidden_states,
             attention_mask=attention_mask,
-            position_ids=position_ids,
+            # position_ids=position_ids,
             past_key_value=past_key_value,
-            output_attentions=output_attentions,
-            use_cache=use_cache,
+            # output_attentions=output_attentions,
+            # use_cache=use_cache,
             cache_position=cache_position,
             position_embeddings=position_embeddings,
-            **kwargs,
+            # **kwargs,
         )
         hidden_states = residual + hidden_states
 
@@ -496,8 +542,8 @@ class OpensciModel(OpensciPreTrainedModel):
         config: OpensciConfig
     """
 
-    def __init__(self, config: OpensciConfig):
-        super().__init__(config)
+    def __init__(self, config: OpensciConfig, **kwargs):
+        super().__init__(config, **kwargs)
         self.padding_idx = config.pad_token_id
         self.vocab_size = config.vocab_size
 
@@ -531,7 +577,7 @@ class OpensciModel(OpensciPreTrainedModel):
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
         cache_position: Optional[torch.LongTensor] = None,
-        **flash_attn_kwargs: Unpack[FlashAttentionKwargs],
+        # **flash_attn_kwargs: Unpack[FlashAttentionKwargs],
     ) -> Union[Tuple, BaseModelOutputWithPast]:
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (
@@ -603,7 +649,7 @@ class OpensciModel(OpensciPreTrainedModel):
                     use_cache=use_cache,
                     cache_position=cache_position,
                     position_embeddings=position_embeddings,
-                    **flash_attn_kwargs,
+                    # **flash_attn_kwargs,
                 )
 
             hidden_states = layer_outputs[0]
@@ -747,15 +793,15 @@ class OpensciModel(OpensciPreTrainedModel):
         return causal_mask
 
 
-class KwargsForCausalLM(FlashAttentionKwargs, TransformersKwargs): ...
+class KwargsForCausalLM(TransformersKwargs): ...
 
 
 class OpensciForCausalLM(OpensciPreTrainedModel, GenerationMixin):
-    _tied_weights_keys = ["lm_head.weight"]
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
     _tp_plan = {"lm_head": "colwise_rep"}
 
-    def __init__(self, config):
-        super().__init__(config)
+    def __init__(self, config, **kwargs):
+        super().__init__(config, **kwargs)
         self.model = OpensciModel(config)
         self.vocab_size = config.vocab_size
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
@@ -891,8 +937,8 @@ class OpensciForCausalLM(OpensciPreTrainedModel, GenerationMixin):
     Opensci_START_DOCSTRING,
 )
 class OpensciForSequenceClassification(OpensciPreTrainedModel):
-    def __init__(self, config):
-        super().__init__(config)
+    def __init__(self, config, **kwargs):
+        super().__init__(config, **kwargs)
         self.num_labels = config.num_labels
         self.model = OpensciModel(config)
         self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
@@ -979,4 +1025,4 @@ class OpensciForSequenceClassification(OpensciPreTrainedModel):
             past_key_values=transformer_outputs.past_key_values,
             hidden_states=transformer_outputs.hidden_states,
             attentions=transformer_outputs.attentions,
-        )
+        )