Update modeling_rwkv_hybrid.py

modeling_rwkv_hybrid.py

@@ -37,63 +37,95 @@ logger = logging.get_logger(__name__)
 
 _CONFIG_FOR_DOC = "RwkvHybridConfig"
 
+
 class RwkvHybridDecoderLayer(nn.Module):
-    def __init__(self, config: RwkvHybridConfig, layer_idx: int, update_v_first, get_v_first):
+    def __init__(self, config: RwkvHybridConfig, layer_idx: int):
         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,
-                                                update_v_first=update_v_first,
-                                                get_v_first=get_v_first)
+                self.self_attn = Rwkv7Attention(
+                    args=config, layer_id=layer_idx)
             elif config.wkv_version == 6:
-                self.self_attn = Rwkv6Attention(args=config, layer_id=layer_idx,
-                                                update_v_first=update_v_first,
-                                                get_v_first=get_v_first)
+                self.self_attn = Rwkv6Attention(
+                    args=config, layer_id=layer_idx)
             else:
                 raise NotImplementedError
-        elif not self.is_rwkv:
-            self.self_attn = Qwen2Attention(config=config, layer_idx=layer_idx)
         else:
-            self.self_attn = None
-            raise NotImplementedError
+            self.self_attn = Qwen2Attention(config=config, layer_idx=layer_idx)
 
         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)   
+            config.hidden_size, eps=config.rms_norm_eps)
+        self.layer_idx = layer_idx
 
-    
     def forward(
         self,
         hidden_states: torch.Tensor,
         attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
         past_key_value: Optional[Cache] = None,
         output_attentions: Optional[bool] = False,
         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
+        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,
         **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
-        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,
-        )
+        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 = residual + hidden_states
 
         # Fully Connected
@@ -106,8 +138,12 @@ 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
@@ -124,6 +160,7 @@ 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,
@@ -146,6 +183,7 @@ 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)`):
@@ -238,11 +276,13 @@ 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, self.update_v_first, self.get_v_first) for layer_idx in range(config.num_hidden_layers)]
+            [RwkvHybridDecoderLayer(config, layer_idx)
+             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)
@@ -266,19 +306,20 @@ 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,
@@ -292,7 +333,12 @@ class RwkvHybridModel(RwkvHybridPreTrainedModel):
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
         cache_position: Optional[torch.LongTensor] = None,
-        **flash_attn_kwargs: Unpack[FlashAttentionKwargs],
+        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,
     ) -> Union[Tuple, BaseModelOutputWithPast]:
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (
@@ -302,7 +348,8 @@ 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(
@@ -317,7 +364,8 @@ 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
             )
@@ -339,6 +387,7 @@ 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,)
 
@@ -353,6 +402,14 @@ 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(
@@ -364,7 +421,14 @@ class RwkvHybridModel(RwkvHybridPreTrainedModel):
                     use_cache=use_cache,
                     cache_position=cache_position,
                     position_embeddings=position_embeddings,
-                    **flash_attn_kwargs,
+                    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)
                 )
 
             hidden_states = layer_outputs[0]
@@ -372,6 +436,14 @@ 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
@@ -402,7 +474,8 @@ 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
@@ -447,7 +520,8 @@ 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
 
@@ -490,16 +564,20 @@ 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
@@ -508,7 +586,9 @@ class RwkvHybridModel(RwkvHybridPreTrainedModel):
         return causal_mask
 
 
-class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs): ...
+class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs):
+    ...
+
 
 class RwkvHybridForCausalLM(RwkvHybridPreTrainedModel, GenerationMixin):
     _tied_weights_keys = ["lm_head.weight"]
@@ -518,7 +598,8 @@ 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()
@@ -548,7 +629,8 @@ 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,
@@ -611,12 +693,15 @@ 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:]
@@ -629,4 +714,3 @@ class RwkvHybridForCausalLM(RwkvHybridPreTrainedModel, GenerationMixin):
             hidden_states=outputs.hidden_states,
             attentions=outputs.attentions,
         )
-