Delete backbone

backbone/bert_layers.py

@@ -1,912 +0,0 @@
-# Copyright 2022 MosaicML Examples authors
-# SPDX-License-Identifier: Apache-2.0
-
-# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
-# Copyright (c) 2018-2021, NVIDIA CORPORATION.  All rights reserved.
-# Copyright (c) 2022, Tri Dao.
-
-import copy
-import logging
-import math
-import warnings
-from typing import List, Optional, Tuple, Union
-
-import torch
-import torch.nn as nn
-from einops import rearrange
-from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present
-from transformers.activations import ACT2FN
-from transformers.modeling_outputs import (MaskedLMOutput,
-                                           SequenceClassifierOutput)
-from transformers.models.bert.modeling_bert import BertPreTrainedModel
-from transformers.modeling_utils import PreTrainedModel
-
-from .bert_padding import (index_first_axis,
-                                            index_put_first_axis, pad_input,
-                                            unpad_input, unpad_input_only)
-
-try:
-    from .flash_attn_triton import flash_attn_qkvpacked_func
-except ImportError as e:
-    flash_attn_qkvpacked_func = None
-
-logger = logging.getLogger(__name__)
-
-
-class BertEmbeddings(nn.Module):
-
-    def __init__(self, config):
-        super().__init__()
-        self.word_embeddings = nn.Embedding(config.vocab_size,
-                                            config.hidden_size,
-                                            padding_idx=config.pad_token_id)
-        # ALiBi doesn't use position embeddings
-        self.token_type_embeddings = nn.Embedding(config.type_vocab_size,
-                                                  config.hidden_size)
-
-        # self.LayerNorm is not snake-cased to stick with TensorFlow model
-        # variable name and be able to load any TensorFlow checkpoint file
-        self.LayerNorm = nn.LayerNorm(config.hidden_size,
-                                      eps=config.layer_norm_eps)
-        self.dropout = nn.Dropout(config.hidden_dropout_prob)
-        self.register_buffer('token_type_ids',
-                             torch.zeros(config.max_position_embeddings,
-                                         dtype=torch.long),
-                             persistent=False)
-
-    def forward(
-        self,
-        input_ids: Optional[torch.LongTensor] = None,
-        token_type_ids: Optional[torch.LongTensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        past_key_values_length: int = 0,
-    ) -> torch.Tensor:
-        if (input_ids is not None) == (inputs_embeds is not None):
-            raise ValueError('Must specify either input_ids or input_embeds!')
-        if input_ids is not None:
-            input_shape = input_ids.size()
-        else:
-            assert inputs_embeds is not None  # just for type checking
-            input_shape = inputs_embeds.size()[:-1]
-
-        seq_length = input_shape[1]
-
-        if position_ids is None:
-            # great! ALiBi
-            pass
-
-        # Setting the token_type_ids to the registered buffer in constructor
-        # where it is all zeros, which usually occurs when it's auto-generated;
-        # registered buffer helps users when tracing the model without passing
-        # token_type_ids, solves issue #5664
-        if token_type_ids is None:
-            if hasattr(self, 'token_type_ids'):
-                assert isinstance(self.token_type_ids, torch.LongTensor)
-                buffered_token_type_ids = self.token_type_ids[:, :seq_length]
-                buffered_token_type_ids_expanded = buffered_token_type_ids.expand(
-                    input_shape[0], seq_length)
-                token_type_ids = buffered_token_type_ids_expanded  # type: ignore
-            else:
-                token_type_ids = torch.zeros(input_shape,  # type: ignore
-                                             dtype=torch.long,
-                                             device=self.word_embeddings.device) # type: ignore  # yapf: disable
-
-        if inputs_embeds is None:
-            inputs_embeds = self.word_embeddings(input_ids)
-        token_type_embeddings = self.token_type_embeddings(token_type_ids)
-
-        embeddings = inputs_embeds + token_type_embeddings
-        # no position embeddings! ALiBi
-        embeddings = self.LayerNorm(embeddings)
-        embeddings = self.dropout(embeddings)
-        return embeddings
-
-
-class BertUnpadSelfAttention(nn.Module):
-
-    def __init__(self, config):
-        super().__init__()
-        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(
-                config, 'embedding_size'):
-            raise ValueError(
-                f'The hidden size ({config.hidden_size}) is not a multiple of the number of attention '
-                f'heads ({config.num_attention_heads})')
-
-        self.num_attention_heads = config.num_attention_heads
-        self.attention_head_size = int(config.hidden_size /
-                                       config.num_attention_heads)
-        self.all_head_size = self.num_attention_heads * self.attention_head_size
-        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
-        self.p_dropout = config.attention_probs_dropout_prob
-        self.Wqkv = nn.Linear(self.all_head_size, 3 * config.hidden_size)
-
-        # Warn if defaulting to pytorch because of import issues
-        if flash_attn_qkvpacked_func is None:
-            warnings.warn(
-                'Unable to import Triton; defaulting MosaicBERT attention implementation to pytorch (this will reduce throughput when using this model).'
-            )
-
-    def forward(self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor,
-                max_seqlen_in_batch: int, indices: torch.Tensor,
-                attn_mask: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
-        """Perform self-attention.
-
-        If dropout is zero, then we can use the Triton kernel, so we do that. However, if not, we send through a standard PyTorch
-        implementation of self-attention.
-
-        The arguments are unpadded, and our implementations of attention require padded arguments,
-        so we first call `pad_input`. Once we compute attention, we re-unpad our outputs for the other layers.
-        The pad/unpad operations add overhead, but not sending pad tokens through ffs saves compute.
-        It is possible to write an unpadded implementation of attention (in Triton and PyTorch), which we will eventually do.
-
-        Args:
-            hidden_states: (total_nnz, dim)
-            cu_seqlens: (batch + 1,)
-            max_seqlen_in_batch: int
-            indices: (total_nnz,)
-            attn_mask: (batch, max_seqlen_in_batch)
-            bias: (batch, heads, max_seqlen_in_batch, max_seqlen_in_batch)
-
-        Returns:
-            attention: (total_nnz, dim)
-        """
-        qkv = self.Wqkv(hidden_states)
-        qkv = pad_input(qkv, indices, cu_seqlens.shape[0] - 1,
-                        max_seqlen_in_batch)  # batch, max_seqlen_in_batch, thd
-        qkv = rearrange(qkv,
-                        'b s (t h d) -> b s t h d',
-                        t=3,
-                        h=self.num_attention_heads)
-        if self.p_dropout or flash_attn_qkvpacked_func is None:
-            # if we have nonzero attention dropout (e.g. during fine-tuning) or no Triton, compute attention in PyTorch
-            q = qkv[:, :, 0, :, :].permute(0, 2, 1, 3)  # b h s d
-            k = qkv[:, :, 1, :, :].permute(0, 2, 3, 1)  # b h d s
-            v = qkv[:, :, 2, :, :].permute(0, 2, 1, 3)  # b h s d
-            attention_scores = torch.matmul(q, k) / math.sqrt(
-                self.attention_head_size)
-            attention_scores = attention_scores + bias
-            attention_probs = nn.functional.softmax(attention_scores, dim=-1)
-            attention_probs = self.dropout(attention_probs)
-            attention = torch.matmul(attention_probs, v).permute(0, 2, 1,
-                                                                 3)  # b s h d
-        else:
-            # Triton implementation only supports 0 attention dropout
-            convert_dtype = qkv.dtype not in [torch.float16, torch.bfloat16]
-            if convert_dtype:
-                # Triton implementation only supports fp16 and bf16
-                orig_dtype = qkv.dtype
-                qkv = qkv.to(torch.float16)
-                bias_dtype = bias.dtype
-                bias = bias.to(torch.float16)
-                attention = flash_attn_qkvpacked_func(qkv, bias)
-                attention = attention.to(orig_dtype)
-                bias = bias.to(bias_dtype)
-            else:
-                attention = flash_attn_qkvpacked_func(qkv, bias)
-
-        # attn_mask is 1 for attend and 0 for don't
-        attention = unpad_input_only(attention, torch.squeeze(attn_mask) == 1)
-        return rearrange(attention, 'nnz h d -> nnz (h d)')
-
-
-# Copy of transformer's library BertSelfOutput that will not be caught by surgery methods looking for HF BERT modules.
-class BertSelfOutput(nn.Module):
-
-    def __init__(self, config):
-        super().__init__()
-        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
-        self.LayerNorm = nn.LayerNorm(config.hidden_size,
-                                      eps=config.layer_norm_eps)
-        self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
-    def forward(self, hidden_states: torch.Tensor,
-                input_tensor: torch.Tensor) -> torch.Tensor:
-        hidden_states = self.dense(hidden_states)
-        hidden_states = self.dropout(hidden_states)
-        hidden_states = self.LayerNorm(hidden_states + input_tensor)
-        return hidden_states
-
-
-class BertUnpadAttention(nn.Module):
-    """Chains attention, Dropout, and LayerNorm for Mosaic BERT."""
-
-    def __init__(self, config):
-        super().__init__()
-        self.self = BertUnpadSelfAttention(config)
-        self.output = BertSelfOutput(config)
-
-    def forward(
-        self,
-        input_tensor: torch.Tensor,
-        cu_seqlens: torch.Tensor,
-        max_s: int,
-        subset_idx: Optional[torch.Tensor] = None,
-        indices: Optional[torch.Tensor] = None,
-        attn_mask: Optional[torch.Tensor] = None,
-        bias: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        """Forward pass for scaled self-attention without padding.
-
-        Arguments:
-            input_tensor: (total_nnz, dim)
-            cu_seqlens: (batch + 1,)
-            max_s: int
-            subset_idx: () set of indices whose values we care about at the end of the layer
-                        (e.g., the masked tokens, if this is the final layer).
-            indices: None or (total_nnz,)
-            attn_mask: None or (batch, max_seqlen_in_batch)
-            bias: None or (batch, heads, max_seqlen_in_batch, max_seqlen_in_batch)
-        """
-        self_output = self.self(input_tensor, cu_seqlens, max_s, indices,
-                                attn_mask, bias)
-        if subset_idx is not None:
-            return self.output(index_first_axis(self_output, subset_idx),
-                               index_first_axis(input_tensor, subset_idx))
-        else:
-            return self.output(self_output, input_tensor)
-
-
-class BertGatedLinearUnitMLP(nn.Module):
-    """Applies the FFN at the end of each Mosaic BERT layer.
-
-    Compared to the default BERT architecture, this block replaces :class:`~transformers.model.bert.modeling_bert.BertIntermediate`
-    and :class:`~transformers.model.bert.modeling_bert.SelfOutput` with a single module that has similar functionality, but
-    introduces Gated Linear Units.
-
-    Note: Mosaic BERT adds parameters in order to implement Gated Linear Units. To keep parameter count consistent with that of a
-    standard Hugging Face BERT, scale down `config.intermediate_size` by 2/3. For example, a Mosaic BERT constructed with
-    `config.intermediate_size=2048` will have the same parameter footprint as its Hugging Face BERT counterpart constructed
-    with the `config.intermediate_size=3072`.
-    However, in most cases it will not be necessary to adjust `config.intermediate_size` since, despite the increased
-    parameter size, Mosaic BERT typically offers a net higher throughput than a Hugging Face BERT built from the same `config`.
-    """
-
-    def __init__(self, config):
-        super().__init__()
-        self.config = config
-        self.gated_layers = nn.Linear(config.hidden_size,
-                                      config.intermediate_size * 2,
-                                      bias=False)
-        self.act = nn.GELU(approximate='none')
-        self.wo = nn.Linear(config.intermediate_size, config.hidden_size)
-        self.dropout = nn.Dropout(config.hidden_dropout_prob)
-        self.layernorm = nn.LayerNorm(config.hidden_size,
-                                      eps=config.layer_norm_eps)
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        """Compute new hidden states from current hidden states.
-
-        Args:
-            hidden_states (torch.Tensor): The (unpadded) hidden states from
-                the attention layer [nnz, dim].
-        """
-        residual_connection = hidden_states
-        # compute the activation
-        hidden_states = self.gated_layers(hidden_states)
-        gated = hidden_states[:, :self.config.intermediate_size]
-        non_gated = hidden_states[:, self.config.intermediate_size:]
-        hidden_states = self.act(gated) * non_gated
-        hidden_states = self.dropout(hidden_states)
-        # multiply by the second matrix
-        hidden_states = self.wo(hidden_states)
-        # add the residual connection and post-LN
-        hidden_states = self.layernorm(hidden_states + residual_connection)
-        return hidden_states
-
-
-class BertLayer(nn.Module):
-    """Composes the Mosaic BERT attention and FFN blocks into a single layer."""
-
-    def __init__(self, config):
-        super(BertLayer, self).__init__()
-        self.attention = BertUnpadAttention(config)
-        self.mlp = BertGatedLinearUnitMLP(config)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        cu_seqlens: torch.Tensor,
-        seqlen: int,
-        subset_idx: Optional[torch.Tensor] = None,
-        indices: Optional[torch.Tensor] = None,
-        attn_mask: Optional[torch.Tensor] = None,
-        bias: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        """Forward pass for a BERT layer, including both attention and MLP.
-
-        Args:
-            hidden_states: (total_nnz, dim)
-            cu_seqlens: (batch + 1,)
-            seqlen: int
-            subset_idx: () set of indices whose values we care about at the end of the layer
-                        (e.g., the masked tokens, if this is the final layer).
-            indices: None or (total_nnz,)
-            attn_mask: None or (batch, max_seqlen_in_batch)
-            bias: None or (batch, heads, max_seqlen_in_batch, max_seqlen_in_batch)
-        """
-        attention_output = self.attention(hidden_states, cu_seqlens, seqlen,
-                                          subset_idx, indices, attn_mask, bias)
-        layer_output = self.mlp(attention_output)
-        return layer_output
-
-
-class BertEncoder(nn.Module):
-    """A stack of BERT layers providing the backbone of Mosaic BERT.
-
-    This module is modeled after the Hugging Face BERT's :class:`~transformers.model.bert.modeling_bert.BertEncoder`,
-    but with substantial modifications to implement unpadding and ALiBi.
-
-    Compared to the analogous Hugging Face BERT module, this module handles unpadding to reduce unnecessary computation
-    at padded tokens, and pre-computes attention biases to implement ALiBi.
-    """
-
-    def __init__(self, config):
-        super().__init__()
-        layer = BertLayer(config)
-        self.layer = nn.ModuleList(
-            [copy.deepcopy(layer) for _ in range(config.num_hidden_layers)])
-
-        self.num_attention_heads = config.num_attention_heads
-
-        # The alibi mask will be dynamically expanded if it is too small for
-        # the input the model receives. But it generally helps to initialize it
-        # to a reasonably large size to help pre-allocate CUDA memory.
-        # The default `alibi_starting_size` is 512.
-        self._current_alibi_size = int(config.alibi_starting_size)
-        self.alibi = torch.zeros(
-            (1, self.num_attention_heads, self._current_alibi_size,
-             self._current_alibi_size))
-        self.rebuild_alibi_tensor(size=config.alibi_starting_size)
-
-    def rebuild_alibi_tensor(self,
-                             size: int,
-                             device: Optional[Union[torch.device, str]] = None):
-        # Alibi
-        # Following https://github.com/ofirpress/attention_with_linear_biases/issues/5 (Implementation 1)
-        # In the causal case, you can exploit the fact that softmax is invariant to a uniform translation
-        # of the logits, which makes the math work out *after* applying causal masking. If no causal masking
-        # will be applied, it is necessary to construct the diagonal mask.
-        n_heads = self.num_attention_heads
-
-        def _get_alibi_head_slopes(n_heads: int) -> List[float]:
-
-            def get_slopes_power_of_2(n_heads: int) -> List[float]:
-                start = (2**(-2**-(math.log2(n_heads) - 3)))
-                ratio = start
-                return [start * ratio**i for i in range(n_heads)]
-
-            # In the paper, they only train models that have 2^a heads for some a. This function
-            # has some good properties that only occur when the input is a power of 2. To
-            # maintain that even when the number of heads is not a power of 2, we use a
-            # workaround.
-            if math.log2(n_heads).is_integer():
-                return get_slopes_power_of_2(n_heads)
-
-            closest_power_of_2 = 2**math.floor(math.log2(n_heads))
-            slopes_a = get_slopes_power_of_2(closest_power_of_2)
-            slopes_b = _get_alibi_head_slopes(2 * closest_power_of_2)
-            slopes_b = slopes_b[0::2][:n_heads - closest_power_of_2]
-            return slopes_a + slopes_b
-
-        context_position = torch.arange(size, device=device)[:, None]
-        memory_position = torch.arange(size, device=device)[None, :]
-        relative_position = torch.abs(memory_position - context_position)
-        # [n_heads, max_token_length, max_token_length]
-        relative_position = relative_position.unsqueeze(0).expand(
-            n_heads, -1, -1)
-        slopes = torch.Tensor(_get_alibi_head_slopes(n_heads)).to(device)
-        alibi = slopes.unsqueeze(1).unsqueeze(1) * -relative_position
-        # [1, n_heads, max_token_length, max_token_length]
-        alibi = alibi.unsqueeze(0)
-        assert alibi.shape == torch.Size([1, n_heads, size, size])
-
-        self._current_alibi_size = size
-        self.alibi = alibi
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: torch.Tensor,
-        output_all_encoded_layers: Optional[bool] = True,
-        subset_mask: Optional[torch.Tensor] = None,
-    ) -> List[torch.Tensor]:
-
-        extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
-        extended_attention_mask = extended_attention_mask.to(
-            dtype=torch.float32)  # fp16 compatibility
-        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
-
-        attention_mask_bool = attention_mask.bool()
-        batch, seqlen = hidden_states.shape[:2]
-        # Unpad inputs and mask. It will remove tokens that are padded.
-        # Assume ntokens is total number of tokens (padded and non-padded)
-        # and ntokens_unpad is total number of non-padded tokens.
-        # Then unpadding performs the following compression of the inputs:
-        # hidden_states[ntokens,hidden] -> hidden_states[ntokens_unpad,hidden]
-        hidden_states, indices, cu_seqlens, _ = unpad_input(
-            hidden_states, attention_mask_bool)
-
-        # Add alibi matrix to extended_attention_mask
-        if self._current_alibi_size < seqlen:
-            # Rebuild the alibi tensor when needed
-            warnings.warn(
-                f'Increasing alibi size from {self._current_alibi_size} to {seqlen}'
-            )
-            self.rebuild_alibi_tensor(size=seqlen, device=hidden_states.device)
-        elif self.alibi.device != hidden_states.device:
-            # Device catch-up
-            self.alibi = self.alibi.to(hidden_states.device)
-        alibi_bias = self.alibi[:, :, :seqlen, :seqlen]
-        attn_bias = extended_attention_mask[:, :, :seqlen, :seqlen]
-        alibi_attn_mask = attn_bias + alibi_bias
-
-        all_encoder_layers = []
-        if subset_mask is None:
-            for layer_module in self.layer:
-                hidden_states = layer_module(hidden_states,
-                                             cu_seqlens,
-                                             seqlen,
-                                             None,
-                                             indices,
-                                             attn_mask=attention_mask,
-                                             bias=alibi_attn_mask)
-                if output_all_encoded_layers:
-                    all_encoder_layers.append(hidden_states)
-            # Pad inputs and mask. It will insert back zero-padded tokens.
-            # Assume ntokens is total number of tokens (padded and non-padded)
-            # and ntokens_unpad is total number of non-padded tokens.
-            # Then padding performs the following de-compression:
-            #     hidden_states[ntokens_unpad,hidden] -> hidden_states[ntokens,hidden]
-            hidden_states = pad_input(hidden_states, indices, batch, seqlen)
-        else:
-            for i in range(len(self.layer) - 1):
-                layer_module = self.layer[i]
-                hidden_states = layer_module(hidden_states,
-                                             cu_seqlens,
-                                             seqlen,
-                                             None,
-                                             indices,
-                                             attn_mask=attention_mask,
-                                             bias=alibi_attn_mask)
-                if output_all_encoded_layers:
-                    all_encoder_layers.append(hidden_states)
-            subset_idx = torch.nonzero(subset_mask[attention_mask_bool],
-                                       as_tuple=False).flatten()
-            hidden_states = self.layer[-1](hidden_states,
-                                           cu_seqlens,
-                                           seqlen,
-                                           subset_idx=subset_idx,
-                                           indices=indices,
-                                           attn_mask=attention_mask,
-                                           bias=alibi_attn_mask)
-
-        if not output_all_encoded_layers:
-            all_encoder_layers.append(hidden_states)
-        return all_encoder_layers
-
-
-class BertPooler(nn.Module):
-
-    def __init__(self, config):
-        super(BertPooler, self).__init__()
-        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
-        self.activation = nn.Tanh()
-
-    def forward(self,
-                hidden_states: torch.Tensor,
-                pool: Optional[bool] = True) -> torch.Tensor:
-        # We "pool" the model by simply taking the hidden state corresponding
-        # to the first token.
-        first_token_tensor = hidden_states[:, 0] if pool else hidden_states
-        pooled_output = self.dense(first_token_tensor)
-        pooled_output = self.activation(pooled_output)
-        return pooled_output
-
-
-class BertPredictionHeadTransform(nn.Module):
-
-    def __init__(self, config):
-        super().__init__()
-        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
-        if isinstance(config.hidden_act, str):
-            self.transform_act_fn = ACT2FN[config.hidden_act]
-        else:
-            self.transform_act_fn = config.hidden_act
-        self.LayerNorm = torch.nn.LayerNorm(config.hidden_size, eps=1e-12)
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        hidden_states = self.dense(hidden_states)
-        hidden_states = self.transform_act_fn(hidden_states)
-        hidden_states = self.LayerNorm(hidden_states)
-        return hidden_states
-
-
-class BertModel(BertPreTrainedModel):
-    """Overall BERT model.
-
-    Args:
-        config: a BertConfig class instance with the configuration to build a new model
-
-    Inputs:
-        `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length]
-            with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts
-            `extract_features.py`, `run_classifier.py` and `run_squad.py`)
-        `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token
-            types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to
-            a `sentence B` token (see BERT paper for more details).
-        `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices
-            selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
-            input sequence length in the current batch. It's the mask that we typically use for attention when
-            a batch has varying length sentences.
-        `output_all_encoded_layers`: boolean which controls the content of the `encoded_layers` output as described below. Default: `True`.
-
-    Outputs: Tuple of (encoded_layers, pooled_output)
-        `encoded_layers`: controlled by `output_all_encoded_layers` argument:
-            - `output_all_encoded_layers=True`: outputs a list of the full sequences of encoded-hidden-states at the end
-                of each attention block (i.e. 12 full sequences for BERT-base, 24 for BERT-large), each
-                encoded-hidden-state is a torch.FloatTensor of size [batch_size, sequence_length, hidden_size],
-            - `output_all_encoded_layers=False`: outputs only the full sequence of hidden-states corresponding
-                to the last attention block of shape [batch_size, sequence_length, hidden_size],
-        `pooled_output`: a torch.FloatTensor of size [batch_size, hidden_size] which is the output of a
-            classifier pretrained on top of the hidden state associated to the first character of the
-            input (`CLS`) to train on the Next-Sentence task (see BERT's paper).
-
-    Example usage:
-    ```python
-    # Already been converted into WordPiece token ids
-    input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
-    input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
-    token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])
-    config = modeling.BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
-        num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)
-    model = BertModel(config=config)
-    all_encoder_layers, pooled_output = model(input_ids, token_type_ids, input_mask)
-    ```
-    """
-
-    def __init__(self, config, add_pooling_layer=True):
-        super(BertModel, self).__init__(config)
-        self.embeddings = BertEmbeddings(config)
-        self.encoder = BertEncoder(config)
-        self.pooler = BertPooler(config) if add_pooling_layer else None
-        self.post_init()
-
-    def get_input_embeddings(self):
-        return self.embeddings.word_embeddings
-
-    def set_input_embeddings(self, value):
-        self.embeddings.word_embeddings = value
-
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        token_type_ids: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.Tensor] = None,
-        output_all_encoded_layers: Optional[bool] = False,
-        masked_tokens_mask: Optional[torch.Tensor] = None,
-        **kwargs
-    ) -> Tuple[Union[List[torch.Tensor], torch.Tensor], Optional[torch.Tensor]]:
-        if attention_mask is None:
-            attention_mask = torch.ones_like(input_ids)
-        if token_type_ids is None:
-            token_type_ids = torch.zeros_like(input_ids)
-
-        embedding_output = self.embeddings(input_ids, token_type_ids,
-                                           position_ids)
-
-        subset_mask = []
-        first_col_mask = []
-
-        if masked_tokens_mask is None:
-            subset_mask = None
-        else:
-            first_col_mask = torch.zeros_like(masked_tokens_mask)
-            first_col_mask[:, 0] = True
-            subset_mask = masked_tokens_mask | first_col_mask
-
-        encoder_outputs = self.encoder(
-            embedding_output,
-            attention_mask,
-            output_all_encoded_layers=output_all_encoded_layers,
-            subset_mask=subset_mask)
-
-        if masked_tokens_mask is None:
-            sequence_output = encoder_outputs[-1]
-            pooled_output = self.pooler(
-                sequence_output) if self.pooler is not None else None
-        else:
-            # TD [2022-03-01]: the indexing here is very tricky.
-            attention_mask_bool = attention_mask.bool()
-            subset_idx = subset_mask[attention_mask_bool]  # type: ignore
-            sequence_output = encoder_outputs[-1][
-                masked_tokens_mask[attention_mask_bool][subset_idx]]
-            if self.pooler is not None:
-                pool_input = encoder_outputs[-1][
-                    first_col_mask[attention_mask_bool][subset_idx]]
-                pooled_output = self.pooler(pool_input, pool=False)
-            else:
-                pooled_output = None
-
-        if not output_all_encoded_layers:
-            encoder_outputs = sequence_output
-
-        if self.pooler is not None:
-            return encoder_outputs, pooled_output
-
-        return encoder_outputs, None
-
-
-###################
-# Bert Heads
-###################
-class BertLMPredictionHead(nn.Module):
-
-    def __init__(self, config, bert_model_embedding_weights):
-        super().__init__()
-        self.transform = BertPredictionHeadTransform(config)
-        # The output weights are the same as the input embeddings, but there is
-        # an output-only bias for each token.
-        self.decoder = nn.Linear(bert_model_embedding_weights.size(1),
-                                 bert_model_embedding_weights.size(0))
-        self.decoder.weight = bert_model_embedding_weights
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        hidden_states = self.transform(hidden_states)
-        hidden_states = self.decoder(hidden_states)
-        return hidden_states
-
-
-class BertOnlyMLMHead(nn.Module):
-
-    def __init__(self, config, bert_model_embedding_weights):
-        super().__init__()
-        self.predictions = BertLMPredictionHead(config,
-                                                bert_model_embedding_weights)
-
-    def forward(self, sequence_output: torch.Tensor) -> torch.Tensor:
-        prediction_scores = self.predictions(sequence_output)
-        return prediction_scores
-
-
-class BertOnlyNSPHead(nn.Module):
-
-    def __init__(self, config):
-        super().__init__()
-        self.seq_relationship = nn.Linear(config.hidden_size, 2)
-
-    def forward(self, pooled_output: torch.Tensor) -> torch.Tensor:
-        seq_relationship_score = self.seq_relationship(pooled_output)
-        return seq_relationship_score
-
-
-
-class BertForMaskedLM(BertPreTrainedModel):
-
-    def __init__(self, config):
-        super().__init__(config)
-
-        if config.is_decoder:
-            warnings.warn(
-                'If you want to use `BertForMaskedLM` make sure `config.is_decoder=False` for '
-                'bi-directional self-attention.')
-
-        self.bert = BertModel(config, add_pooling_layer=False)
-        self.cls = BertOnlyMLMHead(config,
-                                   self.bert.embeddings.word_embeddings.weight)
-
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    def get_output_embeddings(self):
-        return self.cls.predictions.decoder
-
-    def set_output_embeddings(self, new_embeddings):
-        self.cls.predictions.decoder = new_embeddings
-
-    def forward(
-        self,
-        input_ids: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        token_type_ids: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.Tensor] = None,
-        head_mask: Optional[torch.Tensor] = None,
-        inputs_embeds: Optional[torch.Tensor] = None,
-        encoder_hidden_states: Optional[torch.Tensor] = None,
-        encoder_attention_mask: Optional[torch.Tensor] = None,
-        labels: Optional[torch.Tensor] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-    ) -> Union[Tuple[torch.Tensor], MaskedLMOutput]:
-        # labels should be a `torch.LongTensor` of shape
-        # `(batch_size, sequence_length)`. These are used for computing the
-        #  masked language modeling loss.
-        #
-        # Indices should be in `[-100, 0, ..., config.vocab_size]` (see
-        # `input_ids` docstring) Tokens with indices set to `-100` are ignored
-        # (masked), the loss is only computed for the tokens with labels in `[0,
-        # ..., config.vocab_size]`
-        #
-        # Prediction scores are only computed for masked tokens and the (bs,
-        # seqlen) dimensions are flattened
-        if (input_ids is not None) == (inputs_embeds is not None):
-            raise ValueError('Must specify either input_ids or input_embeds!')
-
-        if labels is None:
-            masked_tokens_mask = None
-        else:
-            masked_tokens_mask = labels > 0
-
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        outputs = self.bert(
-            input_ids,
-            attention_mask=attention_mask,
-            token_type_ids=token_type_ids,
-            position_ids=position_ids,
-            head_mask=head_mask,
-            inputs_embeds=inputs_embeds,
-            encoder_hidden_states=encoder_hidden_states,
-            encoder_attention_mask=encoder_attention_mask,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-            masked_tokens_mask=masked_tokens_mask,
-        )
-        
-        sequence_output = outputs[0]
-        prediction_scores = self.cls(sequence_output)
-
-        loss = None
-        if labels is not None:
-            # Compute loss
-            loss_fct = nn.CrossEntropyLoss()
-            masked_token_idx = torch.nonzero(labels.flatten() > 0,
-                                             as_tuple=False).flatten()
-            loss = loss_fct(prediction_scores,
-                            labels.flatten()[masked_token_idx])
-
-            assert input_ids is not None, 'Coding error; please open an issue'
-            batch, seqlen = input_ids.shape[:2]
-            prediction_scores = rearrange(index_put_first_axis(
-                prediction_scores, masked_token_idx, batch * seqlen),
-                                          '(b s) d -> b s d',
-                                          b=batch)
-
-        if not return_dict:
-            output = (prediction_scores,) + outputs[2:]
-            return ((loss,) + output) if loss is not None else output
-
-        return MaskedLMOutput(
-            loss=loss,
-            logits=prediction_scores,
-            hidden_states=outputs[0],
-            attentions=None,
-        )
-
-    def prepare_inputs_for_generation(self, input_ids: torch.Tensor,
-                                      attention_mask: torch.Tensor,
-                                      **model_kwargs):
-        input_shape = input_ids.shape
-        effective_batch_size = input_shape[0]
-
-        #  add a dummy token
-        if self.config.pad_token_id is None:
-            raise ValueError('The PAD token should be defined for generation')
-
-        attention_mask = torch.cat([
-            attention_mask,
-            attention_mask.new_zeros((attention_mask.shape[0], 1))
-        ],
-                                   dim=-1)
-        dummy_token = torch.full((effective_batch_size, 1),
-                                 self.config.pad_token_id,
-                                 dtype=torch.long,
-                                 device=input_ids.device)
-        input_ids = torch.cat([input_ids, dummy_token], dim=1)
-
-        return {'input_ids': input_ids, 'attention_mask': attention_mask}
-
-
-
-class BertForSequenceClassification(BertPreTrainedModel):
-    """Bert Model transformer with a sequence classification/regression head.
-
-    This head is just a linear layer on top of the pooled output. Used for,
-    e.g., GLUE tasks.
-    """
-
-    def __init__(self, config):
-        super().__init__(config)
-        self.num_labels = config.num_labels
-        self.config = config
-
-        self.bert = BertModel(config)
-        classifier_dropout = (config.classifier_dropout
-                              if config.classifier_dropout is not None else
-                              config.hidden_dropout_prob)
-        self.dropout = nn.Dropout(classifier_dropout)
-        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
-
-        # Initialize weights and apply final processing
-        self.post_init()
-
-
-    def forward(
-        self,
-        input_ids: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        token_type_ids: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.Tensor] = None,
-        head_mask: Optional[torch.Tensor] = None,
-        inputs_embeds: Optional[torch.Tensor] = None,
-        labels: Optional[torch.Tensor] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-    ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:
-        # labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
-        # Labels for computing the sequence classification/regression loss.
-        # Indices should be in `[0, ..., config.num_labels - 1]`.
-        # If `config.num_labels == 1` a regression loss is computed
-        # (mean-square loss). If `config.num_labels > 1` a classification loss
-        # is computed (cross-entropy).
-
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        outputs = self.bert(
-            input_ids,
-            attention_mask=attention_mask,
-            token_type_ids=token_type_ids,
-            position_ids=position_ids,
-            head_mask=head_mask,
-            inputs_embeds=inputs_embeds,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-        )
-
-        pooled_output = outputs[1]
-
-        pooled_output = self.dropout(pooled_output)
-        logits = self.classifier(pooled_output)
-
-        loss = None
-        if labels is not None:
-            # Compute loss
-            if self.config.problem_type is None:
-                if self.num_labels == 1:
-                    self.config.problem_type = 'regression'
-                elif self.num_labels > 1 and (labels.dtype == torch.long or
-                                              labels.dtype == torch.int):
-                    self.config.problem_type = 'single_label_classification'
-                else:
-                    self.config.problem_type = 'multi_label_classification'
-
-            if self.config.problem_type == 'regression':
-                loss_fct = nn.MSELoss()
-                if self.num_labels == 1:
-                    loss = loss_fct(logits.squeeze(), labels.squeeze())
-                else:
-                    loss = loss_fct(logits, labels)
-            elif self.config.problem_type == 'single_label_classification':
-                loss_fct = nn.CrossEntropyLoss()
-                loss = loss_fct(logits.view(-1, self.num_labels),
-                                labels.view(-1))
-            elif self.config.problem_type == 'multi_label_classification':
-                loss_fct = nn.BCEWithLogitsLoss()
-                loss = loss_fct(logits, labels)
-
-        if not return_dict:
-            output = (logits,) + outputs[2:]
-            return ((loss,) + output) if loss is not None else output
-
-        return SequenceClassifierOutput(
-            loss=loss,
-            logits=logits,
-            hidden_states=outputs[0],
-            attentions=None,
-        )
-

backbone/bert_padding.py

@@ -1,154 +0,0 @@
-# Copyright 2022 MosaicML Examples authors
-# SPDX-License-Identifier: Apache-2.0
-
-# Adapted from https://github.com/HazyResearch/flash-attention/blob/main/flash_attn/bert_padding.py
-# Which was adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py
-
-
-from typing import Tuple, cast
-
-import torch
-import torch.nn.functional as F
-from einops import rearrange, repeat
-
-
-class IndexFirstAxis(torch.autograd.Function):
-
-    @staticmethod
-    def forward(ctx, input: torch.Tensor,
-                indices: torch.Tensor) -> torch.Tensor:
-        """Get just the values of `input` which are at `indices`.
-
-        Arguments:
-            ctx: the autograd context object
-            input: (b, ...) 2+ dimensional tensor
-            indices: (num_idx) 1D tensor
-        """
-        ctx.save_for_backward(indices)
-        assert input.ndim >= 2
-        ctx.first_axis_dim, other_shape = input.shape[0], input.shape[
-            1:]  # type: ignore
-        second_dim = other_shape.numel(
-        )  # product of sizes of all but first dimension
-        # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing.
-        return torch.gather(
-            rearrange(input, 'b ... -> b (...)'),  # (b, ...) -> (b, second_dim)
-            0,
-            repeat(indices, 'z -> z d',
-                   d=second_dim)  # (indices,) -> (indices, second_dim)
-        ).reshape(-1, *other_shape)  # (num_idx, ...)
-
-    @staticmethod
-    def backward(ctx, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None]:
-        indices, = ctx.saved_tensors
-        assert grad_output.ndim >= 2
-        other_shape = grad_output.shape[1:]
-        grad_output = rearrange(grad_output, 'b ... -> b (...)')
-        grad_input = torch.zeros([ctx.first_axis_dim, grad_output.shape[1]],
-                                 device=grad_output.device,
-                                 dtype=grad_output.dtype)
-        # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing.
-        # grad_input[indices] = grad_output
-        grad_input.scatter_(0,
-                            repeat(indices, 'z -> z d', d=grad_output.shape[1]),
-                            grad_output)
-        return grad_input.reshape(ctx.first_axis_dim, *other_shape), None
-
-
-index_first_axis = IndexFirstAxis.apply
-
-
-class IndexPutFirstAxis(torch.autograd.Function):
-
-    @staticmethod
-    def forward(ctx, values: torch.Tensor, indices: torch.Tensor,
-                first_axis_dim) -> torch.Tensor:
-        ctx.save_for_backward(indices)
-        assert indices.ndim == 1
-        assert values.ndim >= 2
-        output = torch.zeros(first_axis_dim,
-                             *values.shape[1:],
-                             device=values.device,
-                             dtype=values.dtype)
-        output[indices] = values
-        return output
-
-    @staticmethod
-    def backward(ctx,
-                 grad_output: torch.Tensor) -> Tuple[torch.Tensor, None, None]:
-        indices, = ctx.saved_tensors
-        grad_values = grad_output[indices]
-        return grad_values, None, None
-
-
-index_put_first_axis = IndexPutFirstAxis.apply
-
-
-def unpad_input(
-    hidden_states: torch.Tensor,
-    attention_mask: torch.Tensor,
-) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
-    """Remove padding from input sequences.
-
-    Arguments:
-        hidden_states: (batch, seqlen, ...)
-        attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid.
-
-    Returns:
-        hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
-        indices: (total_nnz)
-        cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states.
-        max_seqlen_in_batch: int ()
-    """
-    seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
-    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
-    max_seqlen_in_batch = int(seqlens_in_batch.max().item())
-    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32),
-                       (1, 0))
-    # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the
-    # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim
-    # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to
-    # index with integer indices. Moreover, torch's index is a bit slower than it needs to be,
-    # so we write custom forward and backward to make it a bit faster.
-    hidden_states = cast(
-        torch.Tensor,
-        index_first_axis(rearrange(hidden_states, 'b s ... -> (b s) ...'),
-                         indices))
-    return hidden_states, indices, cu_seqlens, max_seqlen_in_batch
-
-
-def unpad_input_only(
-    hidden_states: torch.Tensor,
-    attention_mask: torch.Tensor,
-) -> torch.Tensor:
-    """Like unpad_input, but only return the unpadded first tensor.
-
-    Save a small amount of overhead.
-
-    Arguments:
-        hidden_states: (batch, seqlen, ...)
-        attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid.
-
-    Returns:
-        hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
-    """
-    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
-    return index_first_axis(rearrange(hidden_states, 'b s ... -> (b s) ...'),
-                            indices)
-
-
-def pad_input(hidden_states: torch.Tensor, indices: torch.Tensor, batch: int,
-              seqlen: int) -> torch.Tensor:
-    """Add padding to sequences.
-
-    Arguments:
-        hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
-        indices: (total_nnz)
-        batch: int batch_size
-        seqlen: int max sequence length
-
-    Returns:
-        hidden_states: (batch, seqlen, ...)
-    """
-    output = index_put_first_axis(hidden_states, indices, batch * seqlen)
-    return rearrange(output, '(b s) ... -> b s ...', b=batch)

backbone/config.json

@@ -1,76 +0,0 @@
-{
-  "_name_or_path": "/data/lb/dnabert2/DNABERT_2-main/finetune/output/dnabert2/v1",
-  "alibi_starting_size": 512,
-  "architectures": [
-    "BertForSequenceClassification"
-  ],
-  "attention_probs_dropout_prob": 0.0,
-  "auto_map": {
-    "AutoConfig": "configuration_bert.BertConfig",
-    "AutoModel": "bert_layers.BertModel",
-    "AutoModelForMaskedLM": "bert_layers.BertForMaskedLM",
-    "AutoModelForSequenceClassification": "bert_layers.BertForSequenceClassification"
-  },
-  "classifier_dropout": null,
-  "gradient_checkpointing": false,
-  "hidden_act": "gelu",
-  "hidden_dropout_prob": 0.1,
-  "hidden_size": 768,
-  "id2label": {
-    "0": "LABEL_0",
-    "1": "LABEL_1",
-    "2": "LABEL_2",
-    "3": "LABEL_3",
-    "4": "LABEL_4",
-    "5": "LABEL_5",
-    "6": "LABEL_6",
-    "7": "LABEL_7",
-    "8": "LABEL_8",
-    "9": "LABEL_9",
-    "10": "LABEL_10",
-    "11": "LABEL_11",
-    "12": "LABEL_12",
-    "13": "LABEL_13",
-    "14": "LABEL_14",
-    "15": "LABEL_15",
-    "16": "LABEL_16",
-    "17": "LABEL_17",
-    "18": "LABEL_18",
-    "19": "LABEL_19"
-  },
-  "initializer_range": 0.02,
-  "intermediate_size": 3072,
-  "label2id": {
-    "LABEL_0": 0,
-    "LABEL_1": 1,
-    "LABEL_10": 10,
-    "LABEL_11": 11,
-    "LABEL_12": 12,
-    "LABEL_13": 13,
-    "LABEL_14": 14,
-    "LABEL_15": 15,
-    "LABEL_16": 16,
-    "LABEL_17": 17,
-    "LABEL_18": 18,
-    "LABEL_19": 19,
-    "LABEL_2": 2,
-    "LABEL_3": 3,
-    "LABEL_4": 4,
-    "LABEL_5": 5,
-    "LABEL_6": 6,
-    "LABEL_7": 7,
-    "LABEL_8": 8,
-    "LABEL_9": 9
-  },
-  "layer_norm_eps": 1e-12,
-  "max_position_embeddings": 512,
-  "num_attention_heads": 12,
-  "num_hidden_layers": 12,
-  "position_embedding_type": "absolute",
-  "problem_type": "single_label_classification",
-  "torch_dtype": "float32",
-  "transformers_version": "4.29.2",
-  "type_vocab_size": 2,
-  "use_cache": true,
-  "vocab_size": 4096
-}

backbone/configuration_bert.py

@@ -1,26 +0,0 @@
-# Copyright 2022 MosaicML Examples authors
-# SPDX-License-Identifier: Apache-2.0
-
-from transformers.configuration_utils import PretrainedConfig
-
-
-class BertConfig(PretrainedConfig):
-
-    def __init__(
-        self,
-        alibi_starting_size: int = 512,
-        attention_probs_dropout_prob: float = 0.0,
-        **kwargs,
-    ):
-        """Configuration class for MosaicBert.
-
-        Args:
-            alibi_starting_size (int): Use `alibi_starting_size` to determine how large of an alibi tensor to
-                create when initializing the model. You should be able to ignore this parameter in most cases.
-                Defaults to 512.
-            attention_probs_dropout_prob (float): By default, turn off attention dropout in Mosaic BERT
-                (otherwise, Flash Attention will be off by default). Defaults to 0.0.
-        """
-        super().__init__(
-            attention_probs_dropout_prob=attention_probs_dropout_prob, **kwargs)
-        self.alibi_starting_size = alibi_starting_size

backbone/flash_attn_triton.py

@@ -1,1112 +0,0 @@
-# Copyright 2022 MosaicML Examples authors
-# SPDX-License-Identifier: Apache-2.0
-
-"""Triton implementation of Flash Attention.
-
-# Copyright (c) 2022, Tri Dao.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-*Experimental* implementation of FlashAttention in Triton.
-We use the FlashAttention implementation from Phil Tillet a starting point.
-https://github.com/openai/triton/blob/master/python/tutorials/06-fused-attention.py
-
-Changes:
-- Implement both causal and non-causal attention.
-- Implement both self-attention and cross-attention.
-- Support arbitrary seqlens (not just multiples of 128), for both forward and backward.
-- Support all head dimensions up to 128 (not just 16, 32, 64, 128), for both forward and backward.
-- Support attention bias.
-- Speed up the forward pass a bit, and only store the LSE instead of m and l.
-- Make the backward for d=128 much faster by reducing register spilling.
-- Optionally parallelize the backward pass across seqlen_k, to deal with the case of
-small batch size * nheads.
-
-Caution:
-- If you plan to use headdim other than 64 and 128, you should test for race conditions
-(due to the Triton compiler), as done in tests/test_flash_attn.py
-"test_flash_attn_triton_race_condition". I've tested and fixed many race conditions
-for different head dimensions (40, 48, 64, 128, 80, 88, 96), but I'm still not 100% confident
-that there are none left for other head dimensions.
-Differences between this Triton version and the CUDA version:
-- Triton version doesn't support dropout.
-- Triton forward is generally faster than CUDA forward.
-- Triton backward is faster than CUDA backward when batch * nheads is small, and when headdim=64.
-It is slightly slower when headdim=128 and batch * nheads is large.
-- Triton version doesn't yet support different sequence lengths in a batch (i.e., RaggedTensor/NestedTensor).
-"""
-
-import math
-
-import torch
-import triton  # type: ignore (reportMissingImports)
-import triton.language as tl  # type: ignore (reportMissingImports)
-from einops import repeat
-
-
-@triton.autotune(
-    configs=[
-        triton.Config({
-            'BLOCK_M': 128,
-            'BLOCK_N': 128
-        },
-                      num_warps=8,
-                      num_stages=1),
-        # This config has a race condition when EVEN_M == False, disabling it for now.
-        # triton.Config({"BLOCK_M": 64, "BLOCK_N": 64}, num_warps=4, num_stages=1),
-    ],
-    key=[
-        'CACHE_KEY_SEQLEN_Q', 'CACHE_KEY_SEQLEN_K', 'BIAS_TYPE', 'IS_CAUSAL',
-        'BLOCK_HEADDIM'
-    ])
-@triton.heuristics({
-    'EVEN_M': lambda args: args['seqlen_q'] % args['BLOCK_M'] == 0,
-    'EVEN_N': lambda args: args['seqlen_k'] % args['BLOCK_N'] == 0,
-    'EVEN_HEADDIM': lambda args: args['headdim'] == args['BLOCK_HEADDIM'],
-})
-@triton.jit
-def _fwd_kernel(
-    Q,
-    K,
-    V,
-    Bias,
-    Out,
-    Lse,
-    TMP,  # NOTE: TMP is a scratchpad buffer to workaround a compiler bug
-    softmax_scale,
-    stride_qb,
-    stride_qh,
-    stride_qm,
-    stride_kb,
-    stride_kh,
-    stride_kn,
-    stride_vb,
-    stride_vh,
-    stride_vn,
-    stride_bb,
-    stride_bh,
-    stride_bm,
-    stride_ob,
-    stride_oh,
-    stride_om,
-    nheads,
-    seqlen_q,
-    seqlen_k,
-    seqlen_q_rounded,
-    headdim,
-    CACHE_KEY_SEQLEN_Q,
-    CACHE_KEY_SEQLEN_K,
-    BIAS_TYPE: tl.constexpr,
-    IS_CAUSAL: tl.constexpr,
-    BLOCK_HEADDIM: tl.constexpr,
-    EVEN_M: tl.constexpr,
-    EVEN_N: tl.constexpr,
-    EVEN_HEADDIM: tl.constexpr,
-    BLOCK_M: tl.constexpr,
-    BLOCK_N: tl.constexpr,
-):
-    start_m = tl.program_id(0)
-    off_hb = tl.program_id(1)
-    off_b = off_hb // nheads
-    off_h = off_hb % nheads
-    # off_b = tl.program_id(1)
-    # off_h = tl.program_id(2)
-    # off_hb = off_b * nheads + off_h
-    # initialize offsets
-    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
-    offs_n = tl.arange(0, BLOCK_N)
-    offs_d = tl.arange(0, BLOCK_HEADDIM)
-    # Initialize pointers to Q, K, V
-    # Adding parenthesis around indexing might use int32 math instead of int64 math?
-    # https://github.com/openai/triton/issues/741
-    # I'm seeing a tiny bit of difference (5-7us)
-    q_ptrs = Q + off_b * stride_qb + off_h * stride_qh + (
-        offs_m[:, None] * stride_qm + offs_d[None, :])
-    k_ptrs = K + off_b * stride_kb + off_h * stride_kh + (
-        offs_n[:, None] * stride_kn + offs_d[None, :])
-    v_ptrs = V + off_b * stride_vb + off_h * stride_vh + (
-        offs_n[:, None] * stride_vn + offs_d[None, :])
-    if BIAS_TYPE == 'vector':
-        b_ptrs = Bias + off_b * stride_bb + off_h * stride_bh + offs_n
-    elif BIAS_TYPE == 'matrix':
-        b_ptrs = Bias + off_b * stride_bb + off_h * stride_bh + (
-            offs_m[:, None] * stride_bm + offs_n[None, :])
-    else:
-        raise ValueError("BIAS_TYPE must be one of {'vector', 'matrix'}")
-    # initialize pointer to m and l
-    t_ptrs = TMP + off_hb * seqlen_q_rounded + offs_m
-    lse_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float('inf')
-    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float('inf')
-    acc_o = tl.zeros([BLOCK_M, BLOCK_HEADDIM], dtype=tl.float32)
-    # load q: it will stay in SRAM throughout
-    # [2022-10-30] TD: Triton bug - in the case of EVEN_M=True and EVEN_N=False, if we just call
-    # tl.load(q_ptrs), we get the wrong output!
-    if EVEN_M & EVEN_N:
-        if EVEN_HEADDIM:
-            q = tl.load(q_ptrs)
-        else:
-            q = tl.load(q_ptrs, mask=offs_d[None, :] < headdim, other=0.0)
-    else:
-        if EVEN_HEADDIM:
-            q = tl.load(q_ptrs, mask=offs_m[:, None] < seqlen_q, other=0.0)
-        else:
-            q = tl.load(q_ptrs,
-                        mask=(offs_m[:, None] < seqlen_q) &
-                        (offs_d[None, :] < headdim),
-                        other=0.0)
-    # loop over k, v and update accumulator
-    end_n = seqlen_k if not IS_CAUSAL else tl.minimum(
-        (start_m + 1) * BLOCK_M, seqlen_k)
-    for start_n in range(0, end_n, BLOCK_N):
-        start_n = tl.multiple_of(start_n, BLOCK_N)
-        # -- compute qk ----
-        if EVEN_N & EVEN_M:  # If we just do "if EVEN_N", there seems to be some race condition
-            if EVEN_HEADDIM:
-                k = tl.load(k_ptrs + start_n * stride_kn)
-            else:
-                k = tl.load(k_ptrs + start_n * stride_kn,
-                            mask=offs_d[None, :] < headdim,
-                            other=0.0)
-        else:
-            if EVEN_HEADDIM:
-                k = tl.load(k_ptrs + start_n * stride_kn,
-                            mask=(start_n + offs_n)[:, None] < seqlen_k,
-                            other=0.0)
-            else:
-                k = tl.load(k_ptrs + start_n * stride_kn,
-                            mask=((start_n + offs_n)[:, None] < seqlen_k) &
-                            (offs_d[None, :] < headdim),
-                            other=0.0)
-        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
-        qk += tl.dot(q, k, trans_b=True)
-        # Trying to combine the two masks seem to make the result wrong
-        if not EVEN_N:  # Need to mask out otherwise the softmax is wrong
-            qk += tl.where((start_n + offs_n)[None, :] < seqlen_k, 0,
-                           float('-inf'))
-        if IS_CAUSAL:
-            qk += tl.where(offs_m[:, None] >= (start_n + offs_n)[None, :], 0,
-                           float('-inf'))
-        if BIAS_TYPE != 'none':
-            if BIAS_TYPE == 'vector':
-                if EVEN_N:
-                    bias = tl.load(b_ptrs + start_n).to(tl.float32)
-                else:
-                    bias = tl.load(b_ptrs + start_n,
-                                   mask=(start_n + offs_n) < seqlen_k,
-                                   other=0.0).to(tl.float32)
-                bias = bias[None, :]
-            elif BIAS_TYPE == 'matrix':
-                if EVEN_M & EVEN_N:
-                    bias = tl.load(b_ptrs + start_n).to(tl.float32)
-                else:
-                    bias = tl.load(b_ptrs + start_n,
-                                   mask=(offs_m[:, None] < seqlen_q) &
-                                   ((start_n + offs_n)[None, :] < seqlen_k),
-                                   other=0.0).to(tl.float32)
-            else:
-                raise ValueError(
-                    "BIAS_TYPE must be one of {'vector', 'matrix'}")
-            # Slightly faster to multiply the softmax_scale in the tl.exp below since the compiler
-            # can then fuse the mult and add into an fma instruction. But if we have bias we need to
-            # to multiply with softmax_scale here.
-            qk = qk * softmax_scale + bias
-            m_ij = tl.maximum(tl.max(qk, 1), lse_i)
-            p = tl.exp(qk - m_ij[:, None])
-        else:
-            m_ij = tl.maximum(tl.max(qk, 1) * softmax_scale, lse_i)
-            p = tl.exp(qk * softmax_scale - m_ij[:, None])
-        l_ij = tl.sum(p, 1)
-
-        # scale acc_o
-        acc_o_scale = tl.exp(m_i - m_ij)
-
-        # # -- update output accumulator --
-        # BUG: have to store and immediately load
-        tl.store(t_ptrs, acc_o_scale)
-        acc_o_scale = tl.load(t_ptrs)
-        acc_o = acc_o * acc_o_scale[:, None]
-        # update acc_o
-        if EVEN_N & EVEN_M:  # If we just do "if EVEN_N", there seems to be some race condition
-            if EVEN_HEADDIM:
-                v = tl.load(v_ptrs + start_n * stride_vn)
-            else:
-                v = tl.load(v_ptrs + start_n * stride_vn,
-                            mask=offs_d[None, :] < headdim,
-                            other=0.0)
-        else:
-            if EVEN_HEADDIM:
-                v = tl.load(v_ptrs + start_n * stride_vn,
-                            mask=(start_n + offs_n)[:, None] < seqlen_k,
-                            other=0.0)
-            else:
-                v = tl.load(v_ptrs + start_n * stride_vn,
-                            mask=((start_n + offs_n)[:, None] < seqlen_k) &
-                            (offs_d[None, :] < headdim),
-                            other=0.0)
-        p = p.to(v.dtype)
-        acc_o += tl.dot(p, v)
-
-        # -- update statistics
-        m_i = m_ij
-        l_i_new = tl.exp(lse_i - m_ij) + l_ij
-        lse_i = m_ij + tl.log(l_i_new)
-
-    o_scale = tl.exp(m_i - lse_i)
-    # BUG: have to store and immediately load
-    tl.store(t_ptrs, o_scale)
-    o_scale = tl.load(t_ptrs)
-    acc_o = acc_o * o_scale[:, None]
-    # rematerialize offsets to save registers
-    start_m = tl.program_id(0)
-    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
-    # write back l and m
-    lse_ptrs = Lse + off_hb * seqlen_q_rounded + offs_m
-    tl.store(lse_ptrs, lse_i)
-    # initialize pointers to output
-    offs_n = tl.arange(0, BLOCK_HEADDIM)
-    out_ptrs = Out + off_b * stride_ob + off_h * stride_oh + (
-        offs_m[:, None] * stride_om + offs_n[None, :])
-    if EVEN_M:
-        if EVEN_HEADDIM:
-            tl.store(out_ptrs, acc_o)
-        else:
-            tl.store(out_ptrs, acc_o, mask=offs_d[None, :] < headdim)
-    else:
-        if EVEN_HEADDIM:
-            tl.store(out_ptrs, acc_o, mask=offs_m[:, None] < seqlen_q)
-        else:
-            tl.store(out_ptrs,
-                     acc_o,
-                     mask=(offs_m[:, None] < seqlen_q) &
-                     (offs_d[None, :] < headdim))
-
-
-@triton.jit
-def _bwd_preprocess_do_o_dot(
-    Out,
-    DO,
-    Delta,
-    stride_ob,
-    stride_oh,
-    stride_om,
-    stride_dob,
-    stride_doh,
-    stride_dom,
-    nheads,
-    seqlen_q,
-    seqlen_q_rounded,
-    headdim,
-    BLOCK_M: tl.constexpr,
-    BLOCK_HEADDIM: tl.constexpr,
-):
-    start_m = tl.program_id(0)
-    off_hb = tl.program_id(1)
-    off_b = off_hb // nheads
-    off_h = off_hb % nheads
-    # initialize offsets
-    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
-    offs_d = tl.arange(0, BLOCK_HEADDIM)
-    # load
-    o = tl.load(Out + off_b * stride_ob + off_h * stride_oh +
-                offs_m[:, None] * stride_om + offs_d[None, :],
-                mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim),
-                other=0.0).to(tl.float32)
-    do = tl.load(DO + off_b * stride_dob + off_h * stride_doh +
-                 offs_m[:, None] * stride_dom + offs_d[None, :],
-                 mask=(offs_m[:, None] < seqlen_q) &
-                 (offs_d[None, :] < headdim),
-                 other=0.0).to(tl.float32)
-    delta = tl.sum(o * do, axis=1)
-    # write-back
-    tl.store(Delta + off_hb * seqlen_q_rounded + offs_m, delta)
-
-
-@triton.jit
-def _bwd_kernel_one_col_block(
-    start_n,
-    Q,
-    K,
-    V,
-    Bias,
-    DO,
-    DQ,
-    DK,
-    DV,
-    LSE,
-    D,
-    softmax_scale,
-    stride_qm,
-    stride_kn,
-    stride_vn,
-    stride_bm,
-    stride_dom,
-    stride_dqm,
-    stride_dkn,
-    stride_dvn,
-    seqlen_q,
-    seqlen_k,
-    headdim,
-    ATOMIC_ADD: tl.constexpr,
-    BIAS_TYPE: tl.constexpr,
-    IS_CAUSAL: tl.constexpr,
-    BLOCK_HEADDIM: tl.constexpr,
-    EVEN_M: tl.constexpr,
-    EVEN_N: tl.constexpr,
-    EVEN_HEADDIM: tl.constexpr,
-    BLOCK_M: tl.constexpr,
-    BLOCK_N: tl.constexpr,
-):
-    # We need to make sure begin_m is a multiple of BLOCK_M (not BLOCK_N)
-    begin_m = 0 if not IS_CAUSAL else ((start_n * BLOCK_N) // BLOCK_M) * BLOCK_M
-    # initialize row/col offsets
-    offs_qm = begin_m + tl.arange(0, BLOCK_M)
-    offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N)
-    offs_m = tl.arange(0, BLOCK_M)
-    offs_d = tl.arange(0, BLOCK_HEADDIM)
-    # initialize pointers to value-like data
-    q_ptrs = Q + (offs_qm[:, None] * stride_qm + offs_d[None, :])
-    k_ptrs = K + (offs_n[:, None] * stride_kn + offs_d[None, :])
-    v_ptrs = V + (offs_n[:, None] * stride_vn + offs_d[None, :])
-    do_ptrs = DO + (offs_qm[:, None] * stride_dom + offs_d[None, :])
-    dq_ptrs = DQ + (offs_qm[:, None] * stride_dqm + offs_d[None, :])
-    if BIAS_TYPE == 'vector':
-        b_ptrs = Bias + offs_n
-    elif BIAS_TYPE == 'matrix':
-        b_ptrs = Bias + (offs_qm[:, None] * stride_bm + offs_n[None, :])
-    else:
-        raise ValueError("BIAS_TYPE must be one of {'vector', 'matrix'}")
-    # initialize dv and dk
-    dv = tl.zeros([BLOCK_N, BLOCK_HEADDIM], dtype=tl.float32)
-    dk = tl.zeros([BLOCK_N, BLOCK_HEADDIM], dtype=tl.float32)
-    # k and v stay in SRAM throughout
-    # [2022-10-30] TD: Same bug as the fwd. In the case of EVEN_N=True and EVEN_M=False,
-    # if we just call tl.load(k_ptrs), we get the wrong output!
-    if EVEN_N & EVEN_M:
-        if EVEN_HEADDIM:
-            k = tl.load(k_ptrs)
-            v = tl.load(v_ptrs)
-        else:
-            k = tl.load(k_ptrs, mask=offs_d[None, :] < headdim, other=0.0)
-            v = tl.load(v_ptrs, mask=offs_d[None, :] < headdim, other=0.0)
-    else:
-        if EVEN_HEADDIM:
-            k = tl.load(k_ptrs, mask=offs_n[:, None] < seqlen_k, other=0.0)
-            v = tl.load(v_ptrs, mask=offs_n[:, None] < seqlen_k, other=0.0)
-        else:
-            k = tl.load(k_ptrs,
-                        mask=(offs_n[:, None] < seqlen_k) &
-                        (offs_d[None, :] < headdim),
-                        other=0.0)
-            v = tl.load(v_ptrs,
-                        mask=(offs_n[:, None] < seqlen_k) &
-                        (offs_d[None, :] < headdim),
-                        other=0.0)
-    # loop over rows
-    num_block_m = tl.cdiv(seqlen_q, BLOCK_M)
-    for start_m in range(begin_m, num_block_m * BLOCK_M, BLOCK_M):
-        start_m = tl.multiple_of(start_m, BLOCK_M)
-        offs_m_curr = start_m + offs_m
-        # load q, k, v, do on-chip
-        # Same bug as below. Otherwise gives wrong result for headdim=40, seqlen=(128, 117)
-        if EVEN_M & EVEN_HEADDIM:
-            q = tl.load(q_ptrs)
-        else:
-            if EVEN_HEADDIM:
-                q = tl.load(q_ptrs,
-                            mask=offs_m_curr[:, None] < seqlen_q,
-                            other=0.0)
-            else:
-                q = tl.load(q_ptrs,
-                            mask=(offs_m_curr[:, None] < seqlen_q) &
-                            (offs_d[None, :] < headdim),
-                            other=0.0)
-        # recompute p = softmax(qk, dim=-1).T
-        qk = tl.dot(q, k, trans_b=True)
-        # Trying to combine the two masks seem to make the result wrong
-        if not EVEN_N:  # Need to mask out otherwise the softmax is wrong
-            qk = tl.where(offs_n[None, :] < seqlen_k, qk, float('-inf'))
-        if IS_CAUSAL:
-            qk = tl.where(offs_m_curr[:, None] >= (offs_n[None, :]), qk,
-                          float('-inf'))
-        if BIAS_TYPE != 'none':
-            if BIAS_TYPE == 'vector':
-                if EVEN_N:
-                    bias = tl.load(b_ptrs).to(tl.float32)
-                else:
-                    bias = tl.load(b_ptrs, mask=offs_n < seqlen_k,
-                                   other=0.0).to(tl.float32)
-                bias = bias[None, :]
-            elif BIAS_TYPE == 'matrix':
-                if EVEN_M & EVEN_N:
-                    bias = tl.load(b_ptrs).to(tl.float32)
-                else:
-                    bias = tl.load(b_ptrs,
-                                   mask=(offs_m_curr[:, None] < seqlen_q) &
-                                   (offs_n[None, :] < seqlen_k),
-                                   other=0.0).to(tl.float32)
-            else:
-                raise ValueError(
-                    "BIAS_TYPE must be one of {'vector', 'matrix'}")
-            qk = qk * softmax_scale + bias
-        # There seems to be a race condition when headdim=48/96, and dq, dk, dv are wrong.
-        # Also wrong for headdim=64.
-        if not (EVEN_M & EVEN_HEADDIM):
-            tl.debug_barrier()
-        lse_i = tl.load(LSE + offs_m_curr)
-        if BIAS_TYPE == 'none':
-            p = tl.exp(qk * softmax_scale - lse_i[:, None])
-        else:
-            p = tl.exp(qk - lse_i[:, None])
-        # compute dv
-        # [2022-10-30] TD: A Triton bug: if EVEN_M=True and EVEN_HEADDIM=False, if we call
-        # do = tl.load(do_ptrs, mask=offs_d[None, :] < headdim, other=0.0), we get wrong outputs
-        # in the case of headdim=48/96, seqlen_q & seqlen_k >= 512. If headdim=40 or seqlen < 512,
-        # the output is correct.
-        if EVEN_M & EVEN_HEADDIM:
-            do = tl.load(do_ptrs)
-        else:
-            # [2022-11-01] TD: Triton bug, there's a race condition if we just use m_mask and not d_mask.
-            do = tl.load(do_ptrs,
-                         mask=(offs_m_curr[:, None] < seqlen_q) &
-                         (offs_d[None, :] < headdim),
-                         other=0.0)
-        # if EVEN_M:
-        #     if EVEN_HEADDIM:
-        #         do = tl.load(do_ptrs)
-        #     else:
-        #         do = tl.load(do_ptrs, mask=offs_d[None, :] < headdim, other=0.0)
-        # else:
-        #     if EVEN_HEADDIM:
-        #         do = tl.load(do_ptrs, mask=offs_m_curr[:, None] < seqlen_q, other=0.0)
-        #     else:
-        #         do = tl.load(do_ptrs, mask=(offs_m_curr[:, None] < seqlen_q)
-        #                                    & (offs_d[None, :] < headdim), other=0.0)
-        dv += tl.dot(p.to(do.dtype), do, trans_a=True)
-        # compute dp = dot(v, do)
-        # There seems to be a race condition when headdim=48/96, and dq, dk are wrong.
-        # Also wrong for headdim=128, seqlen=(108, 256), and ATOMIC_ADD=True
-        # Also wrong for headdim=64, seqlen=(1023, 1024), and ATOMIC_ADD=False
-        if not (EVEN_M & EVEN_HEADDIM):
-            tl.debug_barrier()
-        dp = tl.dot(do, v, trans_b=True)
-        # There's a race condition for headdim=48
-        if not EVEN_HEADDIM:
-            tl.debug_barrier()
-        # compute ds = p * (dp - delta[:, None])
-        # Putting the subtraction after the dp matmul (instead of before) is slightly faster
-        Di = tl.load(D + offs_m_curr)
-        # Converting ds to q.dtype here reduces register pressure and makes it much faster
-        # for BLOCK_HEADDIM=128
-        ds = (p * (dp - Di[:, None]) * softmax_scale).to(q.dtype)
-        # compute dk = dot(ds.T, q)
-        dk += tl.dot(ds, q, trans_a=True)
-        # compute dq
-        if not ATOMIC_ADD:
-            if EVEN_M & EVEN_HEADDIM:  # Race condition if we just do EVEN_M
-                dq = tl.load(dq_ptrs, eviction_policy='evict_last')
-                dq += tl.dot(ds, k)
-                tl.store(dq_ptrs, dq, eviction_policy='evict_last')
-            else:
-                if EVEN_HEADDIM:
-                    dq = tl.load(dq_ptrs,
-                                 mask=offs_m_curr[:, None] < seqlen_q,
-                                 other=0.0,
-                                 eviction_policy='evict_last')
-                    dq += tl.dot(ds, k)
-                    tl.store(dq_ptrs,
-                             dq,
-                             mask=offs_m_curr[:, None] < seqlen_q,
-                             eviction_policy='evict_last')
-                else:
-                    dq = tl.load(dq_ptrs,
-                                 mask=(offs_m_curr[:, None] < seqlen_q) &
-                                 (offs_d[None, :] < headdim),
-                                 other=0.0,
-                                 eviction_policy='evict_last')
-                    dq += tl.dot(ds, k)
-                    tl.store(dq_ptrs,
-                             dq,
-                             mask=(offs_m_curr[:, None] < seqlen_q) &
-                             (offs_d[None, :] < headdim),
-                             eviction_policy='evict_last')
-        else:  # If we're parallelizing across the seqlen_k dimension
-            dq = tl.dot(ds, k)
-            if EVEN_M & EVEN_HEADDIM:  # Race condition if we just do EVEN_M
-                tl.atomic_add(dq_ptrs, dq)
-            else:
-                if EVEN_HEADDIM:
-                    tl.atomic_add(dq_ptrs,
-                                  dq,
-                                  mask=offs_m_curr[:, None] < seqlen_q)
-                else:
-                    tl.atomic_add(dq_ptrs,
-                                  dq,
-                                  mask=(offs_m_curr[:, None] < seqlen_q) &
-                                  (offs_d[None, :] < headdim))
-        # increment pointers
-        dq_ptrs += BLOCK_M * stride_dqm
-        q_ptrs += BLOCK_M * stride_qm
-        do_ptrs += BLOCK_M * stride_dom
-        if BIAS_TYPE == 'matrix':
-            b_ptrs += BLOCK_M * stride_bm
-    # write-back
-    dv_ptrs = DV + (offs_n[:, None] * stride_dvn + offs_d[None, :])
-    dk_ptrs = DK + (offs_n[:, None] * stride_dkn + offs_d[None, :])
-    # [2022-11-01] TD: Same bug. In the case of EVEN_N=True and EVEN_M=False,
-    # if we just call tl.store(dv_ptrs), there's a race condition
-    if EVEN_N & EVEN_M:
-        if EVEN_HEADDIM:
-            tl.store(dv_ptrs, dv)
-            tl.store(dk_ptrs, dk)
-        else:
-            tl.store(dv_ptrs, dv, mask=offs_d[None, :] < headdim)
-            tl.store(dk_ptrs, dk, mask=offs_d[None, :] < headdim)
-    else:
-        if EVEN_HEADDIM:
-            tl.store(dv_ptrs, dv, mask=offs_n[:, None] < seqlen_k)
-            tl.store(dk_ptrs, dk, mask=offs_n[:, None] < seqlen_k)
-        else:
-            tl.store(dv_ptrs,
-                     dv,
-                     mask=(offs_n[:, None] < seqlen_k) &
-                     (offs_d[None, :] < headdim))
-            tl.store(dk_ptrs,
-                     dk,
-                     mask=(offs_n[:, None] < seqlen_k) &
-                     (offs_d[None, :] < headdim))
-
-
-def init_to_zero(name):
-    return lambda nargs: nargs[name].zero_()
-
-
-@triton.autotune(
-    configs=[
-        triton.Config(
-            {
-                'BLOCK_M': 128,
-                'BLOCK_N': 128,
-                'SEQUENCE_PARALLEL': False
-            },
-            num_warps=8,
-            num_stages=1,
-            pre_hook=init_to_zero('DQ')),
-        triton.Config(
-            {
-                'BLOCK_M': 128,
-                'BLOCK_N': 128,
-                'SEQUENCE_PARALLEL': True
-            },
-            num_warps=8,
-            num_stages=1,
-            pre_hook=init_to_zero('DQ')),
-        # Other configs seem to give wrong results when seqlen_q % 128 != 0, disabling them for now
-        # # Kernel is buggy (give wrong result) if we set BLOCK_m=128, BLOCK_n=64, num_warps=*4*
-        # triton.Config({"BLOCK_M": 128, "BLOCK_N": 64, "SEQUENCE_PARALLEL": False}, num_warps=8, num_stages=1, pre_hook=init_to_zero('DQ')),
-        # triton.Config({"BLOCK_M": 128, "BLOCK_N": 64, "SEQUENCE_PARALLEL": True}, num_warps=8, num_stages=1, pre_hook=init_to_zero('DQ')),
-        # triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "SEQUENCE_PARALLEL": False}, num_warps=4, num_stages=1, pre_hook=init_to_zero('DQ')),
-        # triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "SEQUENCE_PARALLEL": True}, num_warps=4, num_stages=1, pre_hook=init_to_zero('DQ')),
-    ],
-    key=[
-        'CACHE_KEY_SEQLEN_Q', 'CACHE_KEY_SEQLEN_K', 'BIAS_TYPE', 'IS_CAUSAL',
-        'BLOCK_HEADDIM'
-    ],
-)
-@triton.heuristics({
-    'EVEN_M': lambda args: args['seqlen_q'] % args['BLOCK_M'] == 0,
-    'EVEN_N': lambda args: args['seqlen_k'] % args['BLOCK_N'] == 0,
-    'EVEN_HEADDIM': lambda args: args['headdim'] == args['BLOCK_HEADDIM'],
-})
-@triton.jit
-def _bwd_kernel(
-    Q,
-    K,
-    V,
-    Bias,
-    DO,
-    DQ,
-    DK,
-    DV,
-    LSE,
-    D,
-    softmax_scale,
-    stride_qb,
-    stride_qh,
-    stride_qm,
-    stride_kb,
-    stride_kh,
-    stride_kn,
-    stride_vb,
-    stride_vh,
-    stride_vn,
-    stride_bb,
-    stride_bh,
-    stride_bm,
-    stride_dob,
-    stride_doh,
-    stride_dom,
-    stride_dqb,
-    stride_dqh,
-    stride_dqm,
-    stride_dkb,
-    stride_dkh,
-    stride_dkn,
-    stride_dvb,
-    stride_dvh,
-    stride_dvn,
-    nheads,
-    seqlen_q,
-    seqlen_k,
-    seqlen_q_rounded,
-    headdim,
-    CACHE_KEY_SEQLEN_Q,
-    CACHE_KEY_SEQLEN_K,
-    BIAS_TYPE: tl.constexpr,
-    IS_CAUSAL: tl.constexpr,
-    BLOCK_HEADDIM: tl.constexpr,
-    SEQUENCE_PARALLEL: tl.constexpr,
-    EVEN_M: tl.constexpr,
-    EVEN_N: tl.constexpr,
-    EVEN_HEADDIM: tl.constexpr,
-    BLOCK_M: tl.constexpr,
-    BLOCK_N: tl.constexpr,
-):
-    off_hb = tl.program_id(1)
-    off_b = off_hb // nheads
-    off_h = off_hb % nheads
-    # offset pointers for batch/head
-    Q += off_b * stride_qb + off_h * stride_qh
-    K += off_b * stride_kb + off_h * stride_kh
-    V += off_b * stride_vb + off_h * stride_vh
-    DO += off_b * stride_dob + off_h * stride_doh
-    DQ += off_b * stride_dqb + off_h * stride_dqh
-    DK += off_b * stride_dkb + off_h * stride_dkh
-    DV += off_b * stride_dvb + off_h * stride_dvh
-    if BIAS_TYPE != 'none':
-        Bias += off_b * stride_bb + off_h * stride_bh
-    # pointer to row-wise quantities in value-like data
-    D += off_hb * seqlen_q_rounded
-    LSE += off_hb * seqlen_q_rounded
-    if not SEQUENCE_PARALLEL:
-        num_block_n = tl.cdiv(seqlen_k, BLOCK_N)
-        for start_n in range(0, num_block_n):
-            _bwd_kernel_one_col_block(start_n,
-                                      Q,
-                                      K,
-                                      V,
-                                      Bias,
-                                      DO,
-                                      DQ,
-                                      DK,
-                                      DV,
-                                      LSE,
-                                      D,
-                                      softmax_scale,
-                                      stride_qm,
-                                      stride_kn,
-                                      stride_vn,
-                                      stride_bm,
-                                      stride_dom,
-                                      stride_dqm,
-                                      stride_dkn,
-                                      stride_dvn,
-                                      seqlen_q,
-                                      seqlen_k,
-                                      headdim,
-                                      ATOMIC_ADD=False,
-                                      BIAS_TYPE=BIAS_TYPE,
-                                      IS_CAUSAL=IS_CAUSAL,
-                                      BLOCK_HEADDIM=BLOCK_HEADDIM,
-                                      EVEN_M=EVEN_M,
-                                      EVEN_N=EVEN_N,
-                                      EVEN_HEADDIM=EVEN_HEADDIM,
-                                      BLOCK_M=BLOCK_M,
-                                      BLOCK_N=BLOCK_N)
-    else:
-        start_n = tl.program_id(0)
-        _bwd_kernel_one_col_block(start_n,
-                                  Q,
-                                  K,
-                                  V,
-                                  Bias,
-                                  DO,
-                                  DQ,
-                                  DK,
-                                  DV,
-                                  LSE,
-                                  D,
-                                  softmax_scale,
-                                  stride_qm,
-                                  stride_kn,
-                                  stride_vn,
-                                  stride_bm,
-                                  stride_dom,
-                                  stride_dqm,
-                                  stride_dkn,
-                                  stride_dvn,
-                                  seqlen_q,
-                                  seqlen_k,
-                                  headdim,
-                                  ATOMIC_ADD=True,
-                                  BIAS_TYPE=BIAS_TYPE,
-                                  IS_CAUSAL=IS_CAUSAL,
-                                  BLOCK_HEADDIM=BLOCK_HEADDIM,
-                                  EVEN_M=EVEN_M,
-                                  EVEN_N=EVEN_N,
-                                  EVEN_HEADDIM=EVEN_HEADDIM,
-                                  BLOCK_M=BLOCK_M,
-                                  BLOCK_N=BLOCK_N)
-
-
-def _flash_attn_forward(q, k, v, bias=None, causal=False, softmax_scale=None):
-    # shape constraints
-    batch, seqlen_q, nheads, d = q.shape
-    _, seqlen_k, _, _ = k.shape
-    assert k.shape == (batch, seqlen_k, nheads, d)
-    assert v.shape == (batch, seqlen_k, nheads, d)
-    assert d <= 128, 'FlashAttention only support head dimensions up to 128'
-    assert q.dtype == k.dtype == v.dtype, 'All tensors must have the same type'
-    assert q.dtype in [torch.float16,
-                       torch.bfloat16], 'Only support fp16 and bf16'
-    assert q.is_cuda and k.is_cuda and v.is_cuda
-    softmax_scale = softmax_scale or 1.0 / math.sqrt(d)
-
-    has_bias = bias is not None
-    bias_type = 'none'
-    if has_bias:
-        assert bias.dtype in [q.dtype, torch.float]
-        assert bias.is_cuda
-        assert bias.dim() == 4
-        if bias.stride(-1) != 1:
-            bias = bias.contiguous()
-        if bias.shape[2:] == (1, seqlen_k):
-            bias_type = 'vector'
-        elif bias.shape[2:] == (seqlen_q, seqlen_k):
-            bias_type = 'matrix'
-        else:
-            raise RuntimeError('Last 2 dimensions of bias must be (1, seqlen_k)'
-                               ' or (seqlen_q, seqlen_k)')
-        if bias.shape[:2] == (1, nheads):
-            bias = repeat(bias, '1 h ... -> b h ...', b=batch)
-        elif bias.shape[:2] == (batch, 1):
-            bias = repeat(bias, 'b 1 ... -> b h ...', h=nheads)
-        elif bias.shape[:2] == (1, 1):
-            bias = repeat(bias, '1 h ... -> b h ...', b=batch)
-            bias = repeat(bias, 'b 1 ... -> b h ...', h=nheads)
-        assert bias.shape[:2] == (
-            batch, nheads
-        ), f'First 2 dimensions of bias must be broadcastible to (batch, nheads) = ({batch, nheads}). Bias has shape: {bias.shape}'
-    assert bias is not None  # for type checking
-    bias_strides = (bias.stride(0), bias.stride(1),
-                    bias.stride(2)) if has_bias else (0, 0, 0)
-
-    seqlen_q_rounded = math.ceil(seqlen_q / 128) * 128
-    lse = torch.empty((batch, nheads, seqlen_q_rounded),
-                      device=q.device,
-                      dtype=torch.float32)
-    tmp = torch.empty((batch, nheads, seqlen_q_rounded),
-                      device=q.device,
-                      dtype=torch.float32)
-    o = torch.empty_like(q)
-
-    BLOCK_HEADDIM = max(triton.next_power_of_2(d), 16)
-    # BLOCK = 128
-    # num_warps = 4 if d <= 64 else 8
-    grid = lambda META: (triton.cdiv(seqlen_q, META['BLOCK_M']), batch * nheads)
-    _fwd_kernel[grid](  # type: ignore
-        q,
-        k,
-        v,
-        bias,
-        o,
-        lse,
-        tmp,
-        softmax_scale,
-        q.stride(0),
-        q.stride(2),
-        q.stride(1),
-        k.stride(0),
-        k.stride(2),
-        k.stride(1),
-        v.stride(0),
-        v.stride(2),
-        v.stride(1),
-        *bias_strides,
-        o.stride(0),
-        o.stride(2),
-        o.stride(1),
-        nheads,
-        seqlen_q,
-        seqlen_k,
-        seqlen_q_rounded,
-        d,
-        seqlen_q // 32,
-        seqlen_k // 32,  # key for triton cache (limit number of compilations)
-        # Can't use kwargs here because triton autotune expects key to be args, not kwargs
-        # IS_CAUSAL=causal, BLOCK_HEADDIM=d,
-        bias_type,
-        causal,
-        BLOCK_HEADDIM,
-        # BLOCK_M=BLOCK, BLOCK_N=BLOCK,
-        # num_warps=num_warps,
-        # num_stages=1,
-    )
-    return o, lse, softmax_scale  # softmax_scale could have been updated
-
-
-def _flash_attn_backward(do,
-                         q,
-                         k,
-                         v,
-                         o,
-                         lse,
-                         dq,
-                         dk,
-                         dv,
-                         bias=None,
-                         causal=False,
-                         softmax_scale=None):
-    # Make sure that the last dimension is contiguous
-    if do.stride(-1) != 1:
-        do = do.contiguous()
-    batch, seqlen_q, nheads, d = q.shape
-    _, seqlen_k, _, _ = k.shape
-    # assert d in {16, 32, 64, 128}
-    assert d <= 128
-    seqlen_q_rounded = math.ceil(seqlen_q / 128) * 128
-    assert lse.shape == (batch, nheads, seqlen_q_rounded)
-    assert q.stride(-1) == k.stride(-1) == v.stride(-1) == o.stride(-1) == 1
-    assert dq.stride(-1) == dk.stride(-1) == dv.stride(-1) == 1
-    softmax_scale = softmax_scale or 1.0 / math.sqrt(d)
-    # dq_accum = torch.zeros_like(q, dtype=torch.float32)
-    dq_accum = torch.empty_like(q, dtype=torch.float32)
-    delta = torch.empty_like(lse)
-    # delta = torch.zeros_like(lse)
-
-    BLOCK_HEADDIM = max(triton.next_power_of_2(d), 16)
-    grid = lambda META: (triton.cdiv(seqlen_q, META['BLOCK_M']), batch * nheads)
-    _bwd_preprocess_do_o_dot[grid](  # type: ignore
-        o,
-        do,
-        delta,
-        o.stride(0),
-        o.stride(2),
-        o.stride(1),
-        do.stride(0),
-        do.stride(2),
-        do.stride(1),
-        nheads,
-        seqlen_q,
-        seqlen_q_rounded,
-        d,
-        BLOCK_M=128,
-        BLOCK_HEADDIM=BLOCK_HEADDIM,
-    )
-
-    has_bias = bias is not None
-    bias_type = 'none'
-    if has_bias:
-        assert bias.dtype in [q.dtype, torch.float]
-        assert bias.is_cuda
-        assert bias.dim() == 4
-        assert bias.stride(-1) == 1
-        if bias.shape[2:] == (1, seqlen_k):
-            bias_type = 'vector'
-        elif bias.shape[2:] == (seqlen_q, seqlen_k):
-            bias_type = 'matrix'
-        else:
-            raise RuntimeError('Last 2 dimensions of bias must be (1, seqlen_k)'
-                               ' or (seqlen_q, seqlen_k)')
-        if bias.shape[:2] == (1, nheads):
-            bias = repeat(bias, '1 h ... -> b h ...', b=batch)
-        elif bias.shape[:2] == (batch, 1):
-            bias = repeat(bias, 'b 1 ... -> b h ...', h=nheads)
-        elif bias.shape[:2] == (1, 1):
-            bias = repeat(bias, '1 h ... -> b h ...', b=batch)
-            bias = repeat(bias, 'b 1 ... -> b h ...', h=nheads)
-        assert bias.shape[:2] == (
-            batch, nheads
-        ), f'First 2 dimensions of bias must be broadcastible to (batch, nheads) = ({batch, nheads}). Bias has shape: {bias.shape}'
-    assert bias is not None  # type checking
-    bias_strides = (bias.stride(0), bias.stride(1),
-                    bias.stride(2)) if has_bias else (0, 0, 0)
-
-    # BLOCK_M = 128
-    # BLOCK_N = 64
-    # num_warps = 4
-    grid = lambda META: (triton.cdiv(seqlen_k, META['BLOCK_N'])
-                         if META['SEQUENCE_PARALLEL'] else 1, batch * nheads)
-    _bwd_kernel[grid](  # type: ignore
-        q,
-        k,
-        v,
-        bias,
-        do,
-        dq_accum,
-        dk,
-        dv,
-        lse,
-        delta,
-        softmax_scale,
-        q.stride(0),
-        q.stride(2),
-        q.stride(1),
-        k.stride(0),
-        k.stride(2),
-        k.stride(1),
-        v.stride(0),
-        v.stride(2),
-        v.stride(1),
-        *bias_strides,
-        do.stride(0),
-        do.stride(2),
-        do.stride(1),
-        dq_accum.stride(0),
-        dq_accum.stride(2),
-        dq_accum.stride(1),
-        dk.stride(0),
-        dk.stride(2),
-        dk.stride(1),
-        dv.stride(0),
-        dv.stride(2),
-        dv.stride(1),
-        nheads,
-        seqlen_q,
-        seqlen_k,
-        seqlen_q_rounded,
-        d,
-        seqlen_q // 32,
-        seqlen_k // 32,  # key for triton cache (limit number of compilations)
-        # Can't use kwargs here because triton autotune expects key to be args, not kwargs
-        # IS_CAUSAL=causal, BLOCK_HEADDIM=d,
-        bias_type,
-        causal,
-        BLOCK_HEADDIM,
-        # SEQUENCE_PARALLEL=False,
-        # BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N,
-        # num_warps=num_warps,
-        # num_stages=1,
-    )
-    dq.copy_(dq_accum)
-
-
-class _FlashAttnQKVPackedFunc(torch.autograd.Function):
-
-    @staticmethod
-    def forward(ctx, qkv, bias=None, causal=False, softmax_scale=None):
-        """Forward pass for packed FlashAttention.
-
-        Args:
-            ctx: autograd context
-            qkv: (batch, seqlen, 3, nheads, headdim)
-            bias: optional, shape broadcastible to (batch, nheads, seqlen, seqlen).
-                For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen).
-                ALiBi mask for non-causal would have shape (1, nheads, seqlen, seqlen)
-            causal (bool): whether to incorporate causal attention masking
-            softmax_scale (float, optional): scale factor for softmax
-        """
-        # Make sure that the last dimension is contiguous
-        if qkv.stride(-1) != 1:
-            qkv = qkv.contiguous()
-        o, lse, ctx.softmax_scale = _flash_attn_forward(
-            qkv[:, :, 0],
-            qkv[:, :, 1],
-            qkv[:, :, 2],
-            bias=bias,
-            causal=causal,
-            softmax_scale=softmax_scale)
-        ctx.save_for_backward(qkv, o, lse, bias)
-        ctx.causal = causal
-        return o
-
-    @staticmethod
-    def backward(ctx, do):
-        qkv, o, lse, bias = ctx.saved_tensors
-        assert not ctx.needs_input_grad[
-            1], 'FlashAttention does not support bias gradient yet'
-        # Triton's autotune causes the Tensor._version to change, and so Pytorch autograd
-        # does a memcpy. To avoid this we run in inference_mode, which doesn't track the version.
-        with torch.inference_mode():
-            dqkv = torch.empty_like(qkv)
-            _flash_attn_backward(do,
-                                 qkv[:, :, 0],
-                                 qkv[:, :, 1],
-                                 qkv[:, :, 2],
-                                 o,
-                                 lse,
-                                 dqkv[:, :, 0],
-                                 dqkv[:, :, 1],
-                                 dqkv[:, :, 2],
-                                 bias=bias,
-                                 causal=ctx.causal,
-                                 softmax_scale=ctx.softmax_scale)
-        return dqkv, None, None, None
-
-
-flash_attn_qkvpacked_func = _FlashAttnQKVPackedFunc.apply
-
-
-class _FlashAttnFunc(torch.autograd.Function):
-
-    @staticmethod
-    def forward(ctx, q, k, v, bias=None, causal=False, softmax_scale=None):
-        """Forward pass for FlashAttention.
-
-        Args:
-            ctx: autograd context
-            q: (batch_size, seqlen_q, nheads, headdim)
-            k: (batch_size, seqlen_k, nheads, headdim)
-            v: (batch_size, seqlen_k, nheads, headdim)
-            bias: optional, shape broadcastible to (batch, nheads, seqlen_q, seqlen_k).
-                For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen_k).
-                ALiBi mask for non-causal would have shape (1, nheads, seqlen_q, seqlen_k)
-            causal (bool): whether to incorporate causal attention masking
-            softmax_scale (float, optional): scale factor for softmax
-        """
-        # Make sure that the last dimension is contiguous
-        q, k, v = [
-            x if x.stride(-1) == 1 else x.contiguous() for x in [q, k, v]
-        ]
-        o, lse, ctx.softmax_scale = _flash_attn_forward(
-            q, k, v, bias=bias, causal=causal, softmax_scale=softmax_scale)
-        ctx.save_for_backward(q, k, v, o, lse, bias)
-        ctx.causal = causal
-        return o
-
-    @staticmethod
-    def backward(ctx, do):
-        q, k, v, o, lse, bias = ctx.saved_tensors
-        assert not ctx.needs_input_grad[
-            3], 'FlashAttention does not support bias gradient yet'
-        # Triton's autotune causes the Tensor._version to change, and so Pytorch autograd
-        # does a memcpy. To avoid this we run in inference_mode, which doesn't track the version.
-        with torch.inference_mode():
-            dq = torch.empty_like(q)
-            dk = torch.empty_like(k)
-            dv = torch.empty_like(v)
-            _flash_attn_backward(do,
-                                 q,
-                                 k,
-                                 v,
-                                 o,
-                                 lse,
-                                 dq,
-                                 dk,
-                                 dv,
-                                 bias=bias,
-                                 causal=ctx.causal,
-                                 softmax_scale=ctx.softmax_scale)
-        return dq, dk, dv, None, None, None
-
-
-flash_attn_func = _FlashAttnFunc.apply

backbone/generation_config.json

@@ -1,5 +0,0 @@
-{
-  "_from_model_config": true,
-  "pad_token_id": 0,
-  "transformers_version": "4.28.0"
-}

backbone/pytorch_model.bin

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:7ff39ec77a484dd01070a41bfd6e95cdd7247bec80fe357ab43a4be33687aeba
-size 468354983

backbone/tokenizer_config.json

@@ -1 +0,0 @@
-{"tokenizer_class": "PreTrainedTokenizerFast", "unk_token": "[UNK]", "cls_token": "[CLS]", "sep_token": "[SEP]", "pad_token": "[PAD]", "mask_token": "[MASK]"}