Repository: oumi-ai/oumi
File: src/oumi/models/layers/ring_attention.py (Python)

Code before:
        softmax_scale,
        causal=causal,
    )

    return attn_output


def new_decoder_forward(
    self,
    hidden_states: torch.Tensor,
    attention_mask: Optional[torch.Tensor] = None,
    position_ids: Optional[torch.LongTensor] = None,
    past_key_value: Optional[transformers_cache_utils.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,  # will become mandatory in v4.46
    **kwargs,
) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:
    """New decoder forward."""
    assert isinstance(
        self.self_attn,
        transformers_models.llama.modeling_llama.LlamaAttention,
    ) or isinstance(
        self.self_attn,
        transformers_models.mistral.modeling_mistral.MistralAttention,
    ), (
        "Please toggle on the Flash Attention 2 implementation "
        "when using zigzag ring attention monkey patch."
    )

    residual = hidden_states

    hidden_states = self.input_layernorm(hidden_states)

    # Self Attention
    hidden_states, self_attn_weights, present_key_value = 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,
        **kwargs,
    )
    hidden_states = residual + hidden_states

    # Fully Connected

Review comment (suggestion):
Add a comment that it was `LlamaFlashAttention2`  originally. Note that if they dropped Flash attention support it'd be a problem for this code (it only works for FA)

Code after:
        softmax_scale,
        causal=causal,
    )

    return attn_output


def new_decoder_forward(
    self,
    hidden_states: torch.Tensor,
    attention_mask: Optional[torch.Tensor] = None,
    position_ids: Optional[torch.LongTensor] = None,
    past_key_value: Optional[transformers_cache_utils.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,  # will become mandatory in v4.46
    **kwargs,
) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:
    """New decoder forward."""
    # Was originally LlamaFlashAttention2, but this was deleted in transformers 4.48.0.
    assert isinstance(
        self.self_attn,
        transformers_models.llama.modeling_llama.LlamaAttention,
        # Ditto but with MistralFlashAttention2.
    ) or isinstance(
        self.self_attn,
        transformers_models.mistral.modeling_mistral.MistralAttention,
    ), (
        "Please toggle on the Flash Attention 2 implementation "
        "when using zigzag ring attention monkey patch."
    )

    residual = hidden_states

    hidden_states = self.input_layernorm(hidden_states)

    # Self Attention
    hidden_states, self_attn_weights, present_key_value = 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,
        **kwargs,
    )
    hidden_states = residual + hidden_states