Upload 4 files

configuration_rwkv_hybrid.py

@@ -15,9 +15,9 @@
 # limitations under the License.
 """RwkvHybrid model configuration"""
 
-from transformers.configuration_utils import PretrainedConfig
-from transformers.modeling_rope_utils import rope_config_validation
-from transformers.utils import logging
+from ...configuration_utils import PretrainedConfig
+from ...modeling_rope_utils import rope_config_validation
+from ...utils import logging
 from typing import Optional, Union, List
 
 
@@ -218,15 +218,17 @@ class RwkvHybridConfig(PretrainedConfig):
             raise NotImplementedError(f"Unsupported wkv_version: {self.wkv_version}, \
                                         wkv_version must be 6 or 7")
 
-        if wkv_layers == "full" or wkv_layers == None:
+        if wkv_layers == "full" or wkv_layers is None:
             self.wkv_layers = list(range(num_hidden_layers))
         elif isinstance(wkv_layers, list):
             if all(isinstance(layer, int) for layer in wkv_layers):
                 self.wkv_layers = wkv_layers
             else:
-                raise ValueError("All elements in wkv_layers must be integers.")
+                raise ValueError(
+                    "All elements in wkv_layers must be integers.")
         else:
-            raise TypeError("wkv_layers must be either 'full', None, or a list of integers.")
+            raise TypeError(
+                "wkv_layers must be either 'full', None, or a list of integers.")
 
         # for backward compatibility
         if num_key_value_heads is None:

modeling_rwkv_hybrid.py

@@ -1,22 +1,23 @@
-from typing import List, Optional, Tuple, Union
+from typing import Callable, List, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
 from transformers.cache_utils import Cache
 
-from ...cache_utils import Cache, StaticCache
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, StaticCache
 from .hybrid_cache import HybridCache
-from ...generation import GenerationMixin
-from ...modeling_attn_mask_utils import AttentionMaskConverter
-from ...modeling_flash_attention_utils import FlashAttentionKwargs
-from ...modeling_utils import PreTrainedModel
+from transformers.generation import GenerationMixin
+from transformers.modeling_attn_mask_utils import AttentionMaskConverter
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 
-from ...modeling_outputs import (
+from transformers.modeling_outputs import (
     BaseModelOutputWithPast,
     CausalLMOutputWithPast,
 )
-from ...processing_utils import Unpack
-from ...utils import (
+from transformers.processing_utils import Unpack
+from transformers.utils import (
     LossKwargs,
     add_start_docstrings,
     add_start_docstrings_to_model_forward,
@@ -27,104 +28,72 @@ import threading
 from .wkv import Rwkv7Attention, Rwkv6Attention
 from .configuration_rwkv_hybrid import RwkvHybridConfig
 
-from ..qwen2.modeling_qwen2 import (Qwen2MLP,
-                                    Qwen2RMSNorm,
-                                    Qwen2RotaryEmbedding,
+from transformers.models.qwen2.modeling_qwen2 import (Qwen2MLP, 
+                                    Qwen2RMSNorm, 
+                                    Qwen2RotaryEmbedding, 
                                     Qwen2Attention)
 
 logger = logging.get_logger(__name__)
 
 _CONFIG_FOR_DOC = "RwkvHybridConfig"
 
-
 class RwkvHybridDecoderLayer(nn.Module):
-    def __init__(self, config: RwkvHybridConfig, layer_idx: int):
+    def __init__(self, config: RwkvHybridConfig, layer_idx: int, update_v_first, get_v_first):
         super().__init__()
         self.hidden_size = config.hidden_size
 
         self.is_rwkv = True if layer_idx in config.wkv_layers else False
         if self.is_rwkv:
             if config.wkv_version == 7:
-                self.self_attn = Rwkv7Attention(
-                    args=config, layer_id=layer_idx)
+                self.self_attn = Rwkv7Attention(args=config, layer_id=layer_idx,
+                                                update_v_first=update_v_first,
+                                                get_v_first=get_v_first)
             elif config.wkv_version == 6:
-                self.self_attn = Rwkv6Attention(
-                    args=config, layer_id=layer_idx)
+                self.self_attn = Rwkv6Attention(args=config, layer_id=layer_idx,
+                                                update_v_first=update_v_first,
+                                                get_v_first=get_v_first)
             else:
                 raise NotImplementedError
-        else:
+        elif not self.is_rwkv:
             self.self_attn = Qwen2Attention(config=config, layer_idx=layer_idx)
+        else:
+            self.self_attn = None
+            raise NotImplementedError
 
         self.mlp = Qwen2MLP(config)
         self.input_layernorm = Qwen2RMSNorm(
             config.hidden_size, eps=config.rms_norm_eps)
         self.post_attention_layernorm = Qwen2RMSNorm(
-            config.hidden_size, eps=config.rms_norm_eps)
-        self.layer_idx = layer_idx
+            config.hidden_size, eps=config.rms_norm_eps)   
 
+    
     def forward(
         self,
         hidden_states: torch.Tensor,
         attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
         past_key_value: Optional[Cache] = None,
         output_attentions: Optional[bool] = False,
         use_cache: Optional[bool] = False,
-        cache_position: Optional[torch.Tensor] = None,
-        position_embeddings: Optional[torch.Tensor] = None,
-        sequence_mask: Optional[torch.Tensor] = None,
-        cu_seq_lens_q: Optional[torch.LongTensor] = None,
-        cu_seq_lens_k: Optional[torch.LongTensor] = None,
-        max_length_q: Optional[int] = None,
-        max_length_k: Optional[int] = None,
-        cu_seqlens: Optional[torch.LongTensor] = None,
-        v_first: Optional[torch.LongTensor] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC
         **kwargs,
     ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
-
-        if sequence_mask is not None:
-            assert len(sequence_mask.shape) == 2, (
-                "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] "
-                "for padding purposes (0 indicating padding). "
-                "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed."
-            )
-            hidden_states = hidden_states.mul(
-                sequence_mask[:, -hidden_states.shape[-2]:, None])
-
         residual = hidden_states
 
         hidden_states = self.input_layernorm(hidden_states)
 
         # RWKV attention
-        if self.is_rwkv:
-            hidden_states, self_attn_weights, v_first = self.self_attn(
-                hidden_states=hidden_states,
-                position_ids=position_ids,
-                past_key_value=past_key_value,
-                output_attentions=output_attentions,
-                use_cache=use_cache,
-                cache_position=cache_position,
-                position_embeddings=position_embeddings,
-                cu_seqlens=cu_seqlens,
-                v_first=v_first,
-                **kwargs
-            )
-        else:
-            hidden_states, self_attn_weights = self.self_attn(
-                hidden_states=hidden_states,
-                attention_mask=attention_mask,
-                position_ids=position_ids,
-                past_key_value=past_key_value,
-                output_attentions=output_attentions,
-                use_cache=use_cache,
-                cache_position=cache_position,
-                position_embeddings=position_embeddings,
-                cu_seq_lens_q=cu_seq_lens_q,
-                cu_seq_lens_k=cu_seq_lens_k,
-                max_length_q=max_length_q,
-                max_length_k=max_length_k,
-                **kwargs
-            )
+        hidden_states, self_attn_weights = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+        )
         hidden_states = residual + hidden_states
 
         # Fully Connected
@@ -137,12 +106,8 @@ class RwkvHybridDecoderLayer(nn.Module):
         if output_attentions:
             outputs += (self_attn_weights,)
 
-        if self.is_rwkv:
-            outputs += (v_first,)
-
         return outputs
 
-
 RWKV_HYBRID_START_DOCSTRING = r"""
     This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
     library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
@@ -159,7 +124,6 @@ RWKV_HYBRID_START_DOCSTRING = r"""
             [`~PreTrainedModel.from_pretrained`] method to load the model weights.
 """
 
-
 @add_start_docstrings(
     "The bare RWKV Hybrid Model outputting raw hidden-states without any specific head on top.",
     RWKV_HYBRID_START_DOCSTRING,
@@ -182,7 +146,6 @@ class RwkvHybridPreTrainedModel(PreTrainedModel):
             if module.padding_idx is not None:
                 module.weight.data[module.padding_idx].zero_()
 
-
 RWKV_HYBRID_INPUTS_DOCSTRING = r"""
     Args:
         input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
@@ -275,13 +238,11 @@ class RwkvHybridModel(RwkvHybridPreTrainedModel):
         self.padding_idx = config.pad_token_id
         self.vocab_size = config.vocab_size
 
-        self.embed_tokens = nn.Embedding(
-            config.vocab_size, config.hidden_size, self.padding_idx)
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
         self.thread_local = threading.local()
         self.thread_local.v_first = None
         self.layers = nn.ModuleList(
-            [RwkvHybridDecoderLayer(config, layer_idx)
-             for layer_idx in range(config.num_hidden_layers)]
+            [RwkvHybridDecoderLayer(config, layer_idx, self.update_v_first, self.get_v_first) for layer_idx in range(config.num_hidden_layers)]
         )
         self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.rotary_emb = Qwen2RotaryEmbedding(config=config)
@@ -305,20 +266,19 @@ class RwkvHybridModel(RwkvHybridPreTrainedModel):
             for layer in self.layers:
                 layer.self_attn.time_mixer.post_init()
 
+    def update_v_first(self, new_v_first):
+        """Callback function to update v_first in HybridModel."""
+        self.thread_local.v_first = new_v_first
+
+    def get_v_first(self):
+        return self.thread_local.v_first
+
     def get_input_embeddings(self):
         return self.embed_tokens
 
     def set_input_embeddings(self, value):
         self.embed_tokens = value
 
-    def get_v_first(self, layer_idx: int, use_cache: bool, past_key_value: HybridCache):
-        if layer_idx == 0:
-            return None
-
-        if use_cache:
-            return past_key_value.get_v_first()
-        return self.v_first
-
     @add_start_docstrings_to_model_forward(RWKV_HYBRID_INPUTS_DOCSTRING)
     def forward(
         self,
@@ -332,12 +292,7 @@ class RwkvHybridModel(RwkvHybridPreTrainedModel):
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
         cache_position: Optional[torch.LongTensor] = None,
-        cu_seq_lens_q: Optional[torch.LongTensor] = None,
-        cu_seq_lens_k: Optional[torch.LongTensor] = None,
-        max_length_q: Optional[int] = None,
-        max_length_k: Optional[int] = None,
-        cu_seqlens: Optional[torch.LongTensor] = None,
-        **kwargs,
+        **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 = (
@@ -347,8 +302,7 @@ class RwkvHybridModel(RwkvHybridPreTrainedModel):
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
         if (input_ids is None) ^ (inputs_embeds is not None):
-            raise ValueError(
-                "You must specify exactly one of input_ids or inputs_embeds")
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
 
         if self.gradient_checkpointing and self.training and use_cache:
             logger.warning_once(
@@ -363,8 +317,7 @@ class RwkvHybridModel(RwkvHybridPreTrainedModel):
             past_key_values = HybridCache()
 
         if cache_position is None:
-            past_seen_tokens = past_key_values.get_seq_length(
-            ) if past_key_values is not None else 0
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
             cache_position = torch.arange(
                 past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
             )
@@ -386,7 +339,6 @@ class RwkvHybridModel(RwkvHybridPreTrainedModel):
         all_self_attns = () if output_attentions else None
 
         for decoder_layer in self.layers[: self.config.num_hidden_layers]:
-            first_rwkv_layer = True
             if output_hidden_states:
                 all_hidden_states += (hidden_states,)
 
@@ -401,14 +353,6 @@ class RwkvHybridModel(RwkvHybridPreTrainedModel):
                     use_cache,
                     cache_position,
                     position_embeddings,
-                    attention_mask,
-                    cu_seq_lens_q,
-                    cu_seq_lens_k,
-                    max_length_q,
-                    max_length_k,
-                    cu_seqlens,
-                    self.get_v_first(decoder_layer.layer_idx,
-                                     use_cache, past_key_values)
                 )
             else:
                 layer_outputs = decoder_layer(
@@ -420,14 +364,7 @@ class RwkvHybridModel(RwkvHybridPreTrainedModel):
                     use_cache=use_cache,
                     cache_position=cache_position,
                     position_embeddings=position_embeddings,
-                    sequence_mask=attention_mask,
-                    cu_seq_lens_q=cu_seq_lens_q,
-                    cu_seq_lens_k=cu_seq_lens_k,
-                    max_length_q=max_length_q,
-                    max_length_k=max_length_k,
-                    cu_seqlens=cu_seqlens,
-                    v_first=self.get_v_first(
-                        decoder_layer.layer_idx, use_cache, past_key_values)
+                    **flash_attn_kwargs,
                 )
 
             hidden_states = layer_outputs[0]
@@ -435,14 +372,6 @@ class RwkvHybridModel(RwkvHybridPreTrainedModel):
             if output_attentions:
                 all_self_attns += (layer_outputs[1],)
 
-            if first_rwkv_layer is True and decoder_layer.is_rwkv:
-                v_first = layer_outputs[-1]
-                if use_cache:
-                    past_key_values.update_v_first(v_first)
-                else:
-                    self.register_buffer('v_first', v_first)
-                first_rwkv_layer = False
-
         hidden_states = self.norm(hidden_states)
 
         # add hidden states from the last decoder layer
@@ -473,8 +402,7 @@ class RwkvHybridModel(RwkvHybridPreTrainedModel):
         # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in
         # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail
         # to infer the attention mask.
-        past_seen_tokens = past_key_values.get_seq_length(
-        ) if past_key_values is not None else 0
+        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
         using_static_cache = isinstance(past_key_values, StaticCache)
 
         # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward
@@ -519,8 +447,7 @@ class RwkvHybridModel(RwkvHybridPreTrainedModel):
             # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
             # Details: https://github.com/pytorch/pytorch/issues/110213
             min_dtype = torch.finfo(dtype).min
-            causal_mask = AttentionMaskConverter._unmask_unattended(
-                causal_mask, min_dtype)
+            causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)
 
         return causal_mask
 
@@ -563,20 +490,16 @@ class RwkvHybridModel(RwkvHybridPreTrainedModel):
         else:
             min_dtype = torch.finfo(dtype).min
             causal_mask = torch.full(
-                (sequence_length,
-                 target_length), fill_value=min_dtype, dtype=dtype, device=device
+                (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device
             )
             if sequence_length != 1:
                 causal_mask = torch.triu(causal_mask, diagonal=1)
-            causal_mask *= torch.arange(target_length,
-                                        device=device) > cache_position.reshape(-1, 1)
-            causal_mask = causal_mask[None, None,
-                                      :, :].expand(batch_size, 1, -1, -1)
+            causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
+            causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
             if attention_mask is not None:
                 causal_mask = causal_mask.clone()  # copy to contiguous memory for in-place edit
                 mask_length = attention_mask.shape[-1]
-                padding_mask = causal_mask[:, :, :,
-                                           :mask_length] + attention_mask[:, None, None, :]
+                padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :]
                 padding_mask = padding_mask == 0
                 causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
                     padding_mask, min_dtype
@@ -585,9 +508,7 @@ class RwkvHybridModel(RwkvHybridPreTrainedModel):
         return causal_mask
 
 
-class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs):
-    ...
-
+class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs): ...
 
 class RwkvHybridForCausalLM(RwkvHybridPreTrainedModel, GenerationMixin):
     _tied_weights_keys = ["lm_head.weight"]
@@ -597,8 +518,7 @@ class RwkvHybridForCausalLM(RwkvHybridPreTrainedModel, GenerationMixin):
         super().__init__(config)
         self.model = RwkvHybridModel(config)
         self.vocab_size = config.vocab_size
-        self.lm_head = nn.Linear(
-            config.hidden_size, config.vocab_size, bias=False)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
 
         # Initialize weights and apply final processing
         self.post_init()
@@ -628,8 +548,7 @@ class RwkvHybridForCausalLM(RwkvHybridPreTrainedModel, GenerationMixin):
         input_ids: torch.LongTensor = None,
         attention_mask: Optional[torch.Tensor] = None,
         position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[Union[Cache,
-                                        List[torch.FloatTensor]]] = None,
+        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
         inputs_embeds: Optional[torch.FloatTensor] = None,
         labels: Optional[torch.LongTensor] = None,
         use_cache: Optional[bool] = None,
@@ -692,15 +611,12 @@ class RwkvHybridForCausalLM(RwkvHybridPreTrainedModel, GenerationMixin):
         )
 
         hidden_states = outputs[0]
-        # Only compute necessary logits,
-        # and do not upcast them to float if we are not computing the loss
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
         logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :])
 
         loss = None
         if labels is not None:
-            loss = self.loss_function(
-                logits=logits, labels=labels,
-                vocab_size=self.config.vocab_size, **kwargs)
+            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
 
         if not return_dict:
             output = (logits,) + outputs[1:]
@@ -713,3 +629,4 @@ class RwkvHybridForCausalLM(RwkvHybridPreTrainedModel, GenerationMixin):
             hidden_states=outputs.hidden_states,
             attentions=outputs.attentions,
         )
+

wkv.py

@@ -279,8 +279,7 @@ class Rwkv_Tmix_x070(nn.Module):
             self.a0 + (xa @ self.a1) @ self.a2
         )  # a is "in-context learning rate"
         if self.args.wkv_has_gate:
-            g_delta = torch.sigmoid(xg @ self.g1) @ self.g2
-            g = 1.0 + g_delta
+            g = torch.sigmoid(xg @ self.g1) @ self.g2 + 1.0
         kk = k * self.k_k
         kk = F.normalize(kk.view(B, T, self.n_head, -1),
                          p=2.0, dim=-1, eps=1e-4 if kk.dtype == torch.float16 else 1e-12).view(B, T, C)