lgq12697 commited on
Commit
5babb4a
1 Parent(s): 9583645

Add DNABERT-2 model for sequence conservation prediction

Browse files
README.md CHANGED
@@ -1,3 +1,62 @@
1
- ---
2
- license: cc-by-nc-sa-4.0
3
- ---
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: cc-by-nc-sa-4.0
3
+ widget:
4
+ - text: ATTTTGGAAATCGGAAATCTTCTGTGTTTATCCGTACCAATCCCTGGATCAGTAGTTGGTACACACATACATAGTTGGATCACAACATATCACGAATGAATTTAGGCTAAAAGAGTTAAATACTTACATTAGGGCCAGGTAGGCCAACAACTATCAGAGAACAACAGCGGAAGACAAAATAATATAAGGGCCCGGTTAACATGCCACAAGCAGTCGACTGGGGAACGAGACCTAGAACAAGACCGCACTCCGATCATCTTGTGGGATACGCAAGCGTACCGACAAGGGCTTCTCTTCAACACTCTCCTAAAAGATATATAAATAGCAAGGGTGAGTACCAACCGTACTCAGCAAGCCACCACAACAACAATGCGTATGATAGAGGGTATTTCAAGGAATGGCTTCAGGTTCTTTTGCATAAAGCTAATTTTACAATTCTTTTCACAAGCCTAAAACCTAGCATAGACTGATCAAATTTTAGTACCAGTGTTCACTTTAAACAACGACGGTTCTGTCCACCATCCATTGTGATCCCAAGGATAGCTTCCCGCCATTGAATCGTCATGGTTTTCTAAGGATGTCCACCTTCCCTCCTCTCGGGAAGTGGCTCCATCAGCATAAAATTCATCATGCAATATCCCATCCCCCACAAGTTAAAAATTTAGAGTCTAGCCAAGTGTAATACATGTCCCGGTGCTCAATAACCGCGAGCACGGCTATTCGAATAGATTTGGTTTACTCACACTGCAGTGGATGTACACTTTACCCGCACTCCGCAACTGCCCAACACATGAGCCTCGTCCGAACACATGGGACGCGTCACGGCAAAGCTTTTCGATAACCTCGCATTGGTAGTACCCGCTCCATGAACTTAAATCCTCATGCACTCTAGGCGTCCATGTTTCTAGCAGTGAGAGGAGTTCTGGCGCTCCCGGGAAAGAGAAGTCTCACACGCATATTAAATTATGGTTCAAGTTAAGTTCTCTCTCTCACACACTCA
5
+ tags:
6
+ - DNA
7
+ - biology
8
+ - genomics
9
+ ---
10
+ # Plant foundation DNA large language models
11
+
12
+ The plant DNA large language models (LLMs) contain a series of foundation models based on different model architectures, which are pre-trained on various plant reference genomes.
13
+ All the models have a comparable model size between 90 MB and 150 MB, BPE tokenizer is used for tokenization and 8000 tokens are included in the vocabulary.
14
+
15
+
16
+ **Developed by:** zhangtaolab
17
+
18
+ ### Model Sources
19
+
20
+ - **Repository:** [Plant DNA LLMs](https://github.com/zhangtaolab/plant_DNA_LLMs)
21
+ - **Manuscript:** [Versatile applications of foundation DNA large language models in plant genomes]()
22
+
23
+ ### Architecture
24
+
25
+ The model is trained based on the zhihan1996/DNABERT-2-117M model with modified tokenizer.
26
+
27
+ This model is fine-tuned for predicting sequence conservation.
28
+
29
+ ### How to use
30
+
31
+ Install the runtime library first:
32
+ ```bash
33
+ pip install transformers
34
+ ```
35
+
36
+ Here is a simple code for inference:
37
+ ```python
38
+ from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
39
+
40
+ model_name = 'dnabert2-conservation'
41
+ # load model and tokenizer
42
+ model = AutoModelForSequenceClassification.from_pretrained(f'zhangtaolab/{model_name}', trust_remote_code=True)
43
+ tokenizer = AutoTokenizer.from_pretrained(f'zhangtaolab/{model_name}', trust_remote_code=True)
44
+
45
+ # inference
46
+ sequences = ['ACATGCTAAATTAGTTGGCAATTTTTTCTCAGGTAGCTGGGCACAATTTGGTAGTCCAGTTGAACAAAATCCATTAGCTTCTTTTAGCAAGTCCCCTGGTTTGGGCCCTGCCAGTCCCATTAATACCAACCATTTGTCTGGATTGGCTGCAATTCTTTCCCCACAAGCAACAACCTCTACCAAGATTGCACCGATTGGCAAGGACCCTGGAAGGGCTGCAAATCAGATGTTTTCTAACTCTGGATCAACACAAGGAGCAGCTTTTCAGCATTCTATATCCTTTCCTGAGCAAAATGTAAAGGCAAGTCCTAGGCCTATATCTACTTTTGGTGAATCAAGTTCTAGTGCATCAAGTATTGGAACACTGTCCGGTCCTCAATTTCTTTGGGGAAGCCCAACTCCTTACTCTGAGCATTCAAACACTTCTGCCTGGTCTTCATCTTCGGTGGGGCTTCCATTTACATCTAGTGTCCAAAGGCAGGGTTTCCCATATACTAGTAATCACAGTCCTTTTCTTGGCTCCCACTCTCATCATCATGTTGGATCTGCTCCATCTGGCCTTCCGCTTGATAGGCATTTTAGCTACTTCCCTGAGTCACCTGAAGCTTCTCTCATGAGCCCGGTTGCATTTGGGAATTTAAATCACGGTGATGGGAATTTTATGATGAACAACATTAGTGCTCGTGCATCTGTAGGAGCCGGTGTTGGTCTTTCTGGAAATACCCCTGAAATTAGTTCACCCAATTTCAGAATGATGTCTCTGCCTAGGCATGGTTCCTTGTTCCATGGAAATAGTTTGTATTCTGGACCTGGAGCAACTAACATTGAGGGATTAGCTGAACGTGGACGAAGTAGACGACCTGAAAATGGTGGGAACCAAATTGATAGTAAGAAGCTGTACCAGCTTGATCTTGACAAAATCGTCTGTGGTGAAGATACAAGGACTACTTTAATGATTAAAAACATTCCTAACAAGTAAGAATAACTAAACATCTATCCT',
47
+ 'GTCGCAAAAATTGGGCCACTTGCAGTTCAATCTGTTTAATCAAAATTGCATGTGTATCAACTTTTTGCCCAATACTAGCTATATCACACCTCAACTCTTTAATGTGTTCATCACTAGTGTCGAACCTCCTCATCATTTTGTCCAACATATCCTCAACTCGCGCCATACTATCTCCACCATCCCTAGGAGTAACTTCACGATTTTGAGGAGGGACATAGGGCCCATTCCTGTCGTTTCTATTAGCATAGTTACTCCTGTTAAAGTTGTTGTCGCGGTTGTAGTTTCCATCACGTACATAATGACTCTCACGGTTGTAGTTACCATAGTTCCGACCTGGGTTCCCTTGAACTTGGCGCCAGTTATCCTGATTTGAGCCTTGGGCGCTTGGTCGGAAACCCCCTGTCTGCTCATTTACTGCATAAGTGTCCTCCGCGTAACATCATTAGGAGGTGGTGGTTTAGCAAAGTAGTTGACTGCATTTATCTTTTCTGCACCCCCTGTGACATTTTTTAGTACCAACCCAAGCTCAGTTCTCATCTGAGACATTTCTTCTCGAATCTCATCTGTGGCTCGGTTGTGAGTGGACTGCACTACGAAGGTGTTTTTCCCTGTATCAAACTTCCTAGTACTCCAAGCTTTGTTATTTCGGGAGATTTTCTCTAGTTTTTCTGCAATCTCAACATAAGTGCATTCTCCATAAGATCCACCTGCTATAGTGTCCAACACCGCTTTATTGTTATCATCCTGTCCCCGATAGAAGTATTCCTTCAGTGACTCATCATCTATACGGTGATTTAGAACACTTCTCAAGAATGAGGTGAATCTATCCCAAGAACTACTAACTAACTCTCCTGGTAGTGCCACAAAGCTGTTCACCCTTTCTTTGTGGTTTAACTTCTTGGAGATCGGATAGTAGCGTGCTAAGAAGACATCCCTTAGTTGGTTCCAAGTGAATATGGAGTTGTATGCGAGCTTAGTGAACCACATTGCAGCCTCTCCC']
48
+ pipe = pipeline('text-classification', model=model, tokenizer=tokenizer,
49
+ trust_remote_code=True, top_k=None)
50
+ results = pipe(sequences)
51
+ print(results)
52
+
53
+ ```
54
+
55
+
56
+ ### Training data
57
+ We use BertForSequenceClassification to fine-tune the model.
58
+ Detailed training procedure can be found in our manuscript.
59
+
60
+
61
+ #### Hardware
62
+ Model was trained on a NVIDIA GTX1080Ti GPU (11 GB).
bert_layers.py ADDED
@@ -0,0 +1,916 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2022 MosaicML Examples authors
2
+ # SPDX-License-Identifier: Apache-2.0
3
+
4
+ # Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
5
+ # Copyright (c) 2018-2021, NVIDIA CORPORATION. All rights reserved.
6
+ # Copyright (c) 2022, Tri Dao.
7
+
8
+ import copy
9
+ import logging
10
+ import math
11
+ import warnings
12
+ from typing import List, Optional, Tuple, Union
13
+
14
+ import torch
15
+ import torch.nn as nn
16
+ from einops import rearrange
17
+ from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present
18
+ from transformers.activations import ACT2FN
19
+ from transformers.modeling_outputs import (MaskedLMOutput,
20
+ SequenceClassifierOutput)
21
+ from transformers.models.bert.modeling_bert import BertPreTrainedModel
22
+
23
+ from .bert_padding import (index_first_axis,
24
+ index_put_first_axis, pad_input,
25
+ unpad_input, unpad_input_only)
26
+
27
+ try:
28
+ from .flash_attn_triton import flash_attn_qkvpacked_func
29
+ except ImportError as e:
30
+ flash_attn_qkvpacked_func = None
31
+
32
+ flash_attn_qkvpacked_func = None
33
+ logger = logging.getLogger(__name__)
34
+
35
+
36
+ class BertEmbeddings(nn.Module):
37
+
38
+ def __init__(self, config):
39
+ super().__init__()
40
+ self.word_embeddings = nn.Embedding(config.vocab_size,
41
+ config.hidden_size,
42
+ padding_idx=config.pad_token_id)
43
+ # ALiBi doesn't use position embeddings
44
+ self.token_type_embeddings = nn.Embedding(config.type_vocab_size,
45
+ config.hidden_size)
46
+
47
+ # self.LayerNorm is not snake-cased to stick with TensorFlow model
48
+ # variable name and be able to load any TensorFlow checkpoint file
49
+ self.LayerNorm = nn.LayerNorm(config.hidden_size,
50
+ eps=config.layer_norm_eps)
51
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
52
+ self.register_buffer('token_type_ids',
53
+ torch.zeros(config.max_position_embeddings,
54
+ dtype=torch.long),
55
+ persistent=False)
56
+
57
+ def forward(
58
+ self,
59
+ input_ids: Optional[torch.LongTensor] = None,
60
+ token_type_ids: Optional[torch.LongTensor] = None,
61
+ position_ids: Optional[torch.LongTensor] = None,
62
+ inputs_embeds: Optional[torch.FloatTensor] = None,
63
+ past_key_values_length: int = 0,
64
+ ) -> torch.Tensor:
65
+ if (input_ids is not None) == (inputs_embeds is not None):
66
+ raise ValueError('Must specify either input_ids or input_embeds!')
67
+ if input_ids is not None:
68
+ input_shape = input_ids.size()
69
+ else:
70
+ assert inputs_embeds is not None # just for type checking
71
+ input_shape = inputs_embeds.size()[:-1]
72
+
73
+ seq_length = input_shape[1]
74
+
75
+ if position_ids is None:
76
+ # great! ALiBi
77
+ pass
78
+
79
+ # Setting the token_type_ids to the registered buffer in constructor
80
+ # where it is all zeros, which usually occurs when it's auto-generated;
81
+ # registered buffer helps users when tracing the model without passing
82
+ # token_type_ids, solves issue #5664
83
+ if token_type_ids is None:
84
+ if hasattr(self, 'token_type_ids'):
85
+ assert isinstance(self.token_type_ids, torch.LongTensor)
86
+ buffered_token_type_ids = self.token_type_ids[:, :seq_length]
87
+ buffered_token_type_ids_expanded = buffered_token_type_ids.expand(
88
+ input_shape[0], seq_length)
89
+ token_type_ids = buffered_token_type_ids_expanded # type: ignore
90
+ else:
91
+ token_type_ids = torch.zeros(input_shape, # type: ignore
92
+ dtype=torch.long,
93
+ device=self.word_embeddings.device) # type: ignore # yapf: disable
94
+
95
+ if inputs_embeds is None:
96
+ inputs_embeds = self.word_embeddings(input_ids)
97
+ token_type_embeddings = self.token_type_embeddings(token_type_ids)
98
+
99
+ embeddings = inputs_embeds + token_type_embeddings
100
+ # no position embeddings! ALiBi
101
+ embeddings = self.LayerNorm(embeddings)
102
+ embeddings = self.dropout(embeddings)
103
+ return embeddings
104
+
105
+
106
+ class BertUnpadSelfAttention(nn.Module):
107
+
108
+ def __init__(self, config):
109
+ super().__init__()
110
+ if config.hidden_size % config.num_attention_heads != 0 and not hasattr(
111
+ config, 'embedding_size'):
112
+ raise ValueError(
113
+ f'The hidden size ({config.hidden_size}) is not a multiple of the number of attention '
114
+ f'heads ({config.num_attention_heads})')
115
+
116
+ self.num_attention_heads = config.num_attention_heads
117
+ self.attention_head_size = int(config.hidden_size /
118
+ config.num_attention_heads)
119
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
120
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
121
+ self.p_dropout = config.attention_probs_dropout_prob
122
+ self.Wqkv = nn.Linear(self.all_head_size, 3 * config.hidden_size)
123
+
124
+ # Warn if defaulting to pytorch because of import issues
125
+ if flash_attn_qkvpacked_func is None:
126
+ warnings.warn(
127
+ 'Unable to import Triton; defaulting MosaicBERT attention implementation to pytorch (this will reduce throughput when using this model).'
128
+ )
129
+
130
+ def forward(self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor,
131
+ max_seqlen_in_batch: int, indices: torch.Tensor,
132
+ attn_mask: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
133
+ """Perform self-attention.
134
+
135
+ If dropout is zero, then we can use the Triton kernel, so we do that. However, if not, we send through a standard PyTorch
136
+ implementation of self-attention.
137
+
138
+ The arguments are unpadded, and our implementations of attention require padded arguments,
139
+ so we first call `pad_input`. Once we compute attention, we re-unpad our outputs for the other layers.
140
+ The pad/unpad operations add overhead, but not sending pad tokens through ffs saves compute.
141
+ It is possible to write an unpadded implementation of attention (in Triton and PyTorch), which we will eventually do.
142
+
143
+ Args:
144
+ hidden_states: (total_nnz, dim)
145
+ cu_seqlens: (batch + 1,)
146
+ max_seqlen_in_batch: int
147
+ indices: (total_nnz,)
148
+ attn_mask: (batch, max_seqlen_in_batch)
149
+ bias: (batch, heads, max_seqlen_in_batch, max_seqlen_in_batch)
150
+
151
+ Returns:
152
+ attention: (total_nnz, dim)
153
+ """
154
+ qkv = self.Wqkv(hidden_states)
155
+ qkv = pad_input(qkv, indices, cu_seqlens.shape[0] - 1,
156
+ max_seqlen_in_batch) # batch, max_seqlen_in_batch, thd
157
+ qkv = rearrange(qkv,
158
+ 'b s (t h d) -> b s t h d',
159
+ t=3,
160
+ h=self.num_attention_heads)
161
+ if self.p_dropout or flash_attn_qkvpacked_func is None:
162
+ # if we have nonzero attention dropout (e.g. during fine-tuning) or no Triton, compute attention in PyTorch
163
+ q = qkv[:, :, 0, :, :].permute(0, 2, 1, 3) # b h s d
164
+ k = qkv[:, :, 1, :, :].permute(0, 2, 3, 1) # b h d s
165
+ v = qkv[:, :, 2, :, :].permute(0, 2, 1, 3) # b h s d
166
+ attention_scores = torch.matmul(q, k) / math.sqrt(
167
+ self.attention_head_size)
168
+ attention_scores = attention_scores + bias
169
+ attention_probs = nn.functional.softmax(attention_scores, dim=-1)
170
+ attention_probs = self.dropout(attention_probs)
171
+ attention = torch.matmul(attention_probs, v).permute(0, 2, 1,
172
+ 3) # b s h d
173
+ else:
174
+ # Triton implementation only supports 0 attention dropout
175
+ convert_dtype = qkv.dtype not in [torch.float16, torch.bfloat16]
176
+ if convert_dtype:
177
+ # Triton implementation only supports fp16 and bf16
178
+ orig_dtype = qkv.dtype
179
+ qkv = qkv.to(torch.float16)
180
+ bias_dtype = bias.dtype
181
+ bias = bias.to(torch.float16)
182
+ attention = flash_attn_qkvpacked_func(qkv, bias)
183
+ attention = attention.to(orig_dtype)
184
+ bias = bias.to(bias_dtype)
185
+ else:
186
+ attention = flash_attn_qkvpacked_func(qkv, bias)
187
+
188
+ # attn_mask is 1 for attend and 0 for don't
189
+ attention = unpad_input_only(attention, torch.squeeze(attn_mask) == 1)
190
+ return rearrange(attention, 'nnz h d -> nnz (h d)')
191
+
192
+
193
+ # Copy of transformer's library BertSelfOutput that will not be caught by surgery methods looking for HF BERT modules.
194
+ class BertSelfOutput(nn.Module):
195
+
196
+ def __init__(self, config):
197
+ super().__init__()
198
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
199
+ self.LayerNorm = nn.LayerNorm(config.hidden_size,
200
+ eps=config.layer_norm_eps)
201
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
202
+
203
+ def forward(self, hidden_states: torch.Tensor,
204
+ input_tensor: torch.Tensor) -> torch.Tensor:
205
+ hidden_states = self.dense(hidden_states)
206
+ hidden_states = self.dropout(hidden_states)
207
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
208
+ return hidden_states
209
+
210
+
211
+ class BertUnpadAttention(nn.Module):
212
+ """Chains attention, Dropout, and LayerNorm for Mosaic BERT."""
213
+
214
+ def __init__(self, config):
215
+ super().__init__()
216
+ self.self = BertUnpadSelfAttention(config)
217
+ self.output = BertSelfOutput(config)
218
+
219
+ def forward(
220
+ self,
221
+ input_tensor: torch.Tensor,
222
+ cu_seqlens: torch.Tensor,
223
+ max_s: int,
224
+ subset_idx: Optional[torch.Tensor] = None,
225
+ indices: Optional[torch.Tensor] = None,
226
+ attn_mask: Optional[torch.Tensor] = None,
227
+ bias: Optional[torch.Tensor] = None,
228
+ ) -> torch.Tensor:
229
+ """Forward pass for scaled self-attention without padding.
230
+
231
+ Arguments:
232
+ input_tensor: (total_nnz, dim)
233
+ cu_seqlens: (batch + 1,)
234
+ max_s: int
235
+ subset_idx: () set of indices whose values we care about at the end of the layer
236
+ (e.g., the masked tokens, if this is the final layer).
237
+ indices: None or (total_nnz,)
238
+ attn_mask: None or (batch, max_seqlen_in_batch)
239
+ bias: None or (batch, heads, max_seqlen_in_batch, max_seqlen_in_batch)
240
+ """
241
+ self_output = self.self(input_tensor, cu_seqlens, max_s, indices,
242
+ attn_mask, bias)
243
+ if subset_idx is not None:
244
+ return self.output(index_first_axis(self_output, subset_idx),
245
+ index_first_axis(input_tensor, subset_idx))
246
+ else:
247
+ return self.output(self_output, input_tensor)
248
+
249
+
250
+ class BertGatedLinearUnitMLP(nn.Module):
251
+ """Applies the FFN at the end of each Mosaic BERT layer.
252
+
253
+ Compared to the default BERT architecture, this block replaces :class:`~transformers.model.bert.modeling_bert.BertIntermediate`
254
+ and :class:`~transformers.model.bert.modeling_bert.SelfOutput` with a single module that has similar functionality, but
255
+ introduces Gated Linear Units.
256
+
257
+ Note: Mosaic BERT adds parameters in order to implement Gated Linear Units. To keep parameter count consistent with that of a
258
+ standard Hugging Face BERT, scale down `config.intermediate_size` by 2/3. For example, a Mosaic BERT constructed with
259
+ `config.intermediate_size=2048` will have the same parameter footprint as its Hugging Face BERT counterpart constructed
260
+ with the `config.intermediate_size=3072`.
261
+ However, in most cases it will not be necessary to adjust `config.intermediate_size` since, despite the increased
262
+ parameter size, Mosaic BERT typically offers a net higher throughput than a Hugging Face BERT built from the same `config`.
263
+ """
264
+
265
+ def __init__(self, config):
266
+ super().__init__()
267
+ self.config = config
268
+ self.gated_layers = nn.Linear(config.hidden_size,
269
+ config.intermediate_size * 2,
270
+ bias=False)
271
+ self.act = nn.GELU(approximate='none')
272
+ self.wo = nn.Linear(config.intermediate_size, config.hidden_size)
273
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
274
+ self.layernorm = nn.LayerNorm(config.hidden_size,
275
+ eps=config.layer_norm_eps)
276
+
277
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
278
+ """Compute new hidden states from current hidden states.
279
+
280
+ Args:
281
+ hidden_states (torch.Tensor): The (unpadded) hidden states from
282
+ the attention layer [nnz, dim].
283
+ """
284
+ residual_connection = hidden_states
285
+ # compute the activation
286
+ hidden_states = self.gated_layers(hidden_states)
287
+ gated = hidden_states[:, :self.config.intermediate_size]
288
+ non_gated = hidden_states[:, self.config.intermediate_size:]
289
+ hidden_states = self.act(gated) * non_gated
290
+ hidden_states = self.dropout(hidden_states)
291
+ # multiply by the second matrix
292
+ hidden_states = self.wo(hidden_states)
293
+ # add the residual connection and post-LN
294
+ hidden_states = self.layernorm(hidden_states + residual_connection)
295
+ return hidden_states
296
+
297
+
298
+ class BertLayer(nn.Module):
299
+ """Composes the Mosaic BERT attention and FFN blocks into a single layer."""
300
+
301
+ def __init__(self, config):
302
+ super(BertLayer, self).__init__()
303
+ self.attention = BertUnpadAttention(config)
304
+ self.mlp = BertGatedLinearUnitMLP(config)
305
+
306
+ def forward(
307
+ self,
308
+ hidden_states: torch.Tensor,
309
+ cu_seqlens: torch.Tensor,
310
+ seqlen: int,
311
+ subset_idx: Optional[torch.Tensor] = None,
312
+ indices: Optional[torch.Tensor] = None,
313
+ attn_mask: Optional[torch.Tensor] = None,
314
+ bias: Optional[torch.Tensor] = None,
315
+ ) -> torch.Tensor:
316
+ """Forward pass for a BERT layer, including both attention and MLP.
317
+
318
+ Args:
319
+ hidden_states: (total_nnz, dim)
320
+ cu_seqlens: (batch + 1,)
321
+ seqlen: int
322
+ subset_idx: () set of indices whose values we care about at the end of the layer
323
+ (e.g., the masked tokens, if this is the final layer).
324
+ indices: None or (total_nnz,)
325
+ attn_mask: None or (batch, max_seqlen_in_batch)
326
+ bias: None or (batch, heads, max_seqlen_in_batch, max_seqlen_in_batch)
327
+ """
328
+ attention_output = self.attention(hidden_states, cu_seqlens, seqlen,
329
+ subset_idx, indices, attn_mask, bias)
330
+ layer_output = self.mlp(attention_output)
331
+ return layer_output
332
+
333
+
334
+ class BertEncoder(nn.Module):
335
+ """A stack of BERT layers providing the backbone of Mosaic BERT.
336
+
337
+ This module is modeled after the Hugging Face BERT's :class:`~transformers.model.bert.modeling_bert.BertEncoder`,
338
+ but with substantial modifications to implement unpadding and ALiBi.
339
+
340
+ Compared to the analogous Hugging Face BERT module, this module handles unpadding to reduce unnecessary computation
341
+ at padded tokens, and pre-computes attention biases to implement ALiBi.
342
+ """
343
+
344
+ def __init__(self, config):
345
+ super().__init__()
346
+ layer = BertLayer(config)
347
+ self.layer = nn.ModuleList(
348
+ [copy.deepcopy(layer) for _ in range(config.num_hidden_layers)])
349
+
350
+ self.num_attention_heads = config.num_attention_heads
351
+
352
+ # The alibi mask will be dynamically expanded if it is too small for
353
+ # the input the model receives. But it generally helps to initialize it
354
+ # to a reasonably large size to help pre-allocate CUDA memory.
355
+ # The default `alibi_starting_size` is 512.
356
+ self._current_alibi_size = int(config.alibi_starting_size)
357
+ self.alibi = torch.zeros(
358
+ (1, self.num_attention_heads, self._current_alibi_size,
359
+ self._current_alibi_size))
360
+ self.rebuild_alibi_tensor(size=config.alibi_starting_size)
361
+
362
+ def rebuild_alibi_tensor(self,
363
+ size: int,
364
+ device: Optional[Union[torch.device, str]] = None):
365
+ # Alibi
366
+ # Following https://github.com/ofirpress/attention_with_linear_biases/issues/5 (Implementation 1)
367
+ # In the causal case, you can exploit the fact that softmax is invariant to a uniform translation
368
+ # of the logits, which makes the math work out *after* applying causal masking. If no causal masking
369
+ # will be applied, it is necessary to construct the diagonal mask.
370
+ n_heads = self.num_attention_heads
371
+
372
+ def _get_alibi_head_slopes(n_heads: int) -> List[float]:
373
+
374
+ def get_slopes_power_of_2(n_heads: int) -> List[float]:
375
+ start = (2**(-2**-(math.log2(n_heads) - 3)))
376
+ ratio = start
377
+ return [start * ratio**i for i in range(n_heads)]
378
+
379
+ # In the paper, they only train models that have 2^a heads for some a. This function
380
+ # has some good properties that only occur when the input is a power of 2. To
381
+ # maintain that even when the number of heads is not a power of 2, we use a
382
+ # workaround.
383
+ if math.log2(n_heads).is_integer():
384
+ return get_slopes_power_of_2(n_heads)
385
+
386
+ closest_power_of_2 = 2**math.floor(math.log2(n_heads))
387
+ slopes_a = get_slopes_power_of_2(closest_power_of_2)
388
+ slopes_b = _get_alibi_head_slopes(2 * closest_power_of_2)
389
+ slopes_b = slopes_b[0::2][:n_heads - closest_power_of_2]
390
+ return slopes_a + slopes_b
391
+
392
+ context_position = torch.arange(size, device=device)[:, None]
393
+ memory_position = torch.arange(size, device=device)[None, :]
394
+ relative_position = torch.abs(memory_position - context_position)
395
+ # [n_heads, max_token_length, max_token_length]
396
+ relative_position = relative_position.unsqueeze(0).expand(
397
+ n_heads, -1, -1)
398
+ slopes = torch.Tensor(_get_alibi_head_slopes(n_heads)).to(device)
399
+ alibi = slopes.unsqueeze(1).unsqueeze(1) * -relative_position
400
+ # [1, n_heads, max_token_length, max_token_length]
401
+ alibi = alibi.unsqueeze(0)
402
+ assert alibi.shape == torch.Size([1, n_heads, size, size])
403
+
404
+ self._current_alibi_size = size
405
+ self.alibi = alibi
406
+
407
+ def forward(
408
+ self,
409
+ hidden_states: torch.Tensor,
410
+ attention_mask: torch.Tensor,
411
+ output_all_encoded_layers: Optional[bool] = True,
412
+ subset_mask: Optional[torch.Tensor] = None,
413
+ ) -> List[torch.Tensor]:
414
+
415
+ extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
416
+ extended_attention_mask = extended_attention_mask.to(
417
+ dtype=torch.float32) # fp16 compatibility
418
+ extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
419
+
420
+ attention_mask_bool = attention_mask.bool()
421
+ batch, seqlen = hidden_states.shape[:2]
422
+ # Unpad inputs and mask. It will remove tokens that are padded.
423
+ # Assume ntokens is total number of tokens (padded and non-padded)
424
+ # and ntokens_unpad is total number of non-padded tokens.
425
+ # Then unpadding performs the following compression of the inputs:
426
+ # hidden_states[ntokens,hidden] -> hidden_states[ntokens_unpad,hidden]
427
+ hidden_states, indices, cu_seqlens, _ = unpad_input(
428
+ hidden_states, attention_mask_bool)
429
+
430
+ # Add alibi matrix to extended_attention_mask
431
+ if self._current_alibi_size < seqlen:
432
+ # Rebuild the alibi tensor when needed
433
+ warnings.warn(
434
+ f'Increasing alibi size from {self._current_alibi_size} to {seqlen}'
435
+ )
436
+ self.rebuild_alibi_tensor(size=seqlen, device=hidden_states.device)
437
+ elif self.alibi.device != hidden_states.device:
438
+ # Device catch-up
439
+ self.alibi = self.alibi.to(hidden_states.device)
440
+ alibi_bias = self.alibi[:, :, :seqlen, :seqlen]
441
+ attn_bias = extended_attention_mask[:, :, :seqlen, :seqlen]
442
+ alibi_attn_mask = attn_bias + alibi_bias
443
+
444
+ all_encoder_layers = []
445
+ if subset_mask is None:
446
+ for layer_module in self.layer:
447
+ hidden_states = layer_module(hidden_states,
448
+ cu_seqlens,
449
+ seqlen,
450
+ None,
451
+ indices,
452
+ attn_mask=attention_mask,
453
+ bias=alibi_attn_mask)
454
+ if output_all_encoded_layers:
455
+ all_encoder_layers.append(hidden_states)
456
+ # Pad inputs and mask. It will insert back zero-padded tokens.
457
+ # Assume ntokens is total number of tokens (padded and non-padded)
458
+ # and ntokens_unpad is total number of non-padded tokens.
459
+ # Then padding performs the following de-compression:
460
+ # hidden_states[ntokens_unpad,hidden] -> hidden_states[ntokens,hidden]
461
+ hidden_states = pad_input(hidden_states, indices, batch, seqlen)
462
+ else:
463
+ for i in range(len(self.layer) - 1):
464
+ layer_module = self.layer[i]
465
+ hidden_states = layer_module(hidden_states,
466
+ cu_seqlens,
467
+ seqlen,
468
+ None,
469
+ indices,
470
+ attn_mask=attention_mask,
471
+ bias=alibi_attn_mask)
472
+ if output_all_encoded_layers:
473
+ all_encoder_layers.append(hidden_states)
474
+ subset_idx = torch.nonzero(subset_mask[attention_mask_bool],
475
+ as_tuple=False).flatten()
476
+ hidden_states = self.layer[-1](hidden_states,
477
+ cu_seqlens,
478
+ seqlen,
479
+ subset_idx=subset_idx,
480
+ indices=indices,
481
+ attn_mask=attention_mask,
482
+ bias=alibi_attn_mask)
483
+
484
+ if not output_all_encoded_layers:
485
+ all_encoder_layers.append(hidden_states)
486
+ return all_encoder_layers
487
+
488
+
489
+ class BertPooler(nn.Module):
490
+
491
+ def __init__(self, config):
492
+ super(BertPooler, self).__init__()
493
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
494
+ self.activation = nn.Tanh()
495
+
496
+ def forward(self,
497
+ hidden_states: torch.Tensor,
498
+ pool: Optional[bool] = True) -> torch.Tensor:
499
+ # We "pool" the model by simply taking the hidden state corresponding
500
+ # to the first token.
501
+ first_token_tensor = hidden_states[:, 0] if pool else hidden_states
502
+ pooled_output = self.dense(first_token_tensor)
503
+ pooled_output = self.activation(pooled_output)
504
+ return pooled_output
505
+
506
+
507
+ class BertPredictionHeadTransform(nn.Module):
508
+
509
+ def __init__(self, config):
510
+ super().__init__()
511
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
512
+ if isinstance(config.hidden_act, str):
513
+ self.transform_act_fn = ACT2FN[config.hidden_act]
514
+ else:
515
+ self.transform_act_fn = config.hidden_act
516
+ self.LayerNorm = torch.nn.LayerNorm(config.hidden_size, eps=1e-12)
517
+
518
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
519
+ hidden_states = self.dense(hidden_states)
520
+ hidden_states = self.transform_act_fn(hidden_states)
521
+ hidden_states = self.LayerNorm(hidden_states)
522
+ return hidden_states
523
+
524
+
525
+ class BertModel(BertPreTrainedModel):
526
+ """Overall BERT model.
527
+
528
+ Args:
529
+ config: a BertConfig class instance with the configuration to build a new model
530
+
531
+ Inputs:
532
+ `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length]
533
+ with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts
534
+ `extract_features.py`, `run_classifier.py` and `run_squad.py`)
535
+ `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token
536
+ types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to
537
+ a `sentence B` token (see BERT paper for more details).
538
+ `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices
539
+ selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
540
+ input sequence length in the current batch. It's the mask that we typically use for attention when
541
+ a batch has varying length sentences.
542
+ `output_all_encoded_layers`: boolean which controls the content of the `encoded_layers` output as described below. Default: `True`.
543
+
544
+ Outputs: Tuple of (encoded_layers, pooled_output)
545
+ `encoded_layers`: controlled by `output_all_encoded_layers` argument:
546
+ - `output_all_encoded_layers=True`: outputs a list of the full sequences of encoded-hidden-states at the end
547
+ of each attention block (i.e. 12 full sequences for BERT-base, 24 for BERT-large), each
548
+ encoded-hidden-state is a torch.FloatTensor of size [batch_size, sequence_length, hidden_size],
549
+ - `output_all_encoded_layers=False`: outputs only the full sequence of hidden-states corresponding
550
+ to the last attention block of shape [batch_size, sequence_length, hidden_size],
551
+ `pooled_output`: a torch.FloatTensor of size [batch_size, hidden_size] which is the output of a
552
+ classifier pretrained on top of the hidden state associated to the first character of the
553
+ input (`CLS`) to train on the Next-Sentence task (see BERT's paper).
554
+
555
+ Example usage:
556
+ ```python
557
+ # Already been converted into WordPiece token ids
558
+ input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
559
+ input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
560
+ token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])
561
+ config = modeling.BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
562
+ num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)
563
+ model = BertModel(config=config)
564
+ all_encoder_layers, pooled_output = model(input_ids, token_type_ids, input_mask)
565
+ ```
566
+ """
567
+
568
+ def __init__(self, config, add_pooling_layer=True):
569
+ super(BertModel, self).__init__(config)
570
+ self.embeddings = BertEmbeddings(config)
571
+ self.encoder = BertEncoder(config)
572
+ self.pooler = BertPooler(config) if add_pooling_layer else None
573
+ self.post_init()
574
+
575
+ def get_input_embeddings(self):
576
+ return self.embeddings.word_embeddings
577
+
578
+ def set_input_embeddings(self, value):
579
+ self.embeddings.word_embeddings = value
580
+
581
+ def forward(
582
+ self,
583
+ input_ids: torch.Tensor,
584
+ token_type_ids: Optional[torch.Tensor] = None,
585
+ attention_mask: Optional[torch.Tensor] = None,
586
+ position_ids: Optional[torch.Tensor] = None,
587
+ output_all_encoded_layers: Optional[bool] = False,
588
+ masked_tokens_mask: Optional[torch.Tensor] = None,
589
+ **kwargs
590
+ ) -> Tuple[Union[List[torch.Tensor], torch.Tensor], Optional[torch.Tensor]]:
591
+ if attention_mask is None:
592
+ attention_mask = torch.ones_like(input_ids)
593
+ if token_type_ids is None:
594
+ token_type_ids = torch.zeros_like(input_ids)
595
+
596
+ embedding_output = self.embeddings(input_ids, token_type_ids,
597
+ position_ids)
598
+
599
+ subset_mask = []
600
+ first_col_mask = []
601
+
602
+ if masked_tokens_mask is None:
603
+ subset_mask = None
604
+ else:
605
+ first_col_mask = torch.zeros_like(masked_tokens_mask)
606
+ first_col_mask[:, 0] = True
607
+ subset_mask = masked_tokens_mask | first_col_mask
608
+
609
+ encoder_outputs = self.encoder(
610
+ embedding_output,
611
+ attention_mask,
612
+ output_all_encoded_layers=output_all_encoded_layers,
613
+ subset_mask=subset_mask)
614
+
615
+ if masked_tokens_mask is None:
616
+ sequence_output = encoder_outputs[-1]
617
+ pooled_output = self.pooler(
618
+ sequence_output) if self.pooler is not None else None
619
+ else:
620
+ # TD [2022-03-01]: the indexing here is very tricky.
621
+ attention_mask_bool = attention_mask.bool()
622
+ subset_idx = subset_mask[attention_mask_bool] # type: ignore
623
+ sequence_output = encoder_outputs[-1][
624
+ masked_tokens_mask[attention_mask_bool][subset_idx]]
625
+ if self.pooler is not None:
626
+ pool_input = encoder_outputs[-1][
627
+ first_col_mask[attention_mask_bool][subset_idx]]
628
+ pooled_output = self.pooler(pool_input, pool=False)
629
+ else:
630
+ pooled_output = None
631
+
632
+ if not output_all_encoded_layers:
633
+ encoder_outputs = sequence_output
634
+
635
+ if self.pooler is not None:
636
+ return encoder_outputs, pooled_output
637
+
638
+ return encoder_outputs, None
639
+
640
+
641
+ ###################
642
+ # Bert Heads
643
+ ###################
644
+ class BertLMPredictionHead(nn.Module):
645
+
646
+ def __init__(self, config, bert_model_embedding_weights):
647
+ super().__init__()
648
+ self.transform = BertPredictionHeadTransform(config)
649
+ # The output weights are the same as the input embeddings, but there is
650
+ # an output-only bias for each token.
651
+ self.decoder = nn.Linear(bert_model_embedding_weights.size(1),
652
+ bert_model_embedding_weights.size(0))
653
+ self.decoder.weight = bert_model_embedding_weights
654
+
655
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
656
+ hidden_states = self.transform(hidden_states)
657
+ hidden_states = self.decoder(hidden_states)
658
+ return hidden_states
659
+
660
+
661
+ class BertOnlyMLMHead(nn.Module):
662
+
663
+ def __init__(self, config, bert_model_embedding_weights):
664
+ super().__init__()
665
+ self.predictions = BertLMPredictionHead(config,
666
+ bert_model_embedding_weights)
667
+
668
+ def forward(self, sequence_output: torch.Tensor) -> torch.Tensor:
669
+ prediction_scores = self.predictions(sequence_output)
670
+ return prediction_scores
671
+
672
+
673
+ class BertOnlyNSPHead(nn.Module):
674
+
675
+ def __init__(self, config):
676
+ super().__init__()
677
+ self.seq_relationship = nn.Linear(config.hidden_size, 2)
678
+
679
+ def forward(self, pooled_output: torch.Tensor) -> torch.Tensor:
680
+ seq_relationship_score = self.seq_relationship(pooled_output)
681
+ return seq_relationship_score
682
+
683
+
684
+
685
+ class BertForMaskedLM(BertPreTrainedModel):
686
+
687
+ def __init__(self, config):
688
+ super().__init__(config)
689
+
690
+ if config.is_decoder:
691
+ warnings.warn(
692
+ 'If you want to use `BertForMaskedLM` make sure `config.is_decoder=False` for '
693
+ 'bi-directional self-attention.')
694
+
695
+ self.bert = BertModel(config, add_pooling_layer=False)
696
+ self.cls = BertOnlyMLMHead(config,
697
+ self.bert.embeddings.word_embeddings.weight)
698
+
699
+ # Initialize weights and apply final processing
700
+ self.post_init()
701
+
702
+ def get_output_embeddings(self):
703
+ return self.cls.predictions.decoder
704
+
705
+ def set_output_embeddings(self, new_embeddings):
706
+ self.cls.predictions.decoder = new_embeddings
707
+
708
+ def forward(
709
+ self,
710
+ input_ids: Optional[torch.Tensor] = None,
711
+ attention_mask: Optional[torch.Tensor] = None,
712
+ token_type_ids: Optional[torch.Tensor] = None,
713
+ position_ids: Optional[torch.Tensor] = None,
714
+ head_mask: Optional[torch.Tensor] = None,
715
+ inputs_embeds: Optional[torch.Tensor] = None,
716
+ encoder_hidden_states: Optional[torch.Tensor] = None,
717
+ encoder_attention_mask: Optional[torch.Tensor] = None,
718
+ labels: Optional[torch.Tensor] = None,
719
+ output_attentions: Optional[bool] = None,
720
+ output_hidden_states: Optional[bool] = None,
721
+ return_dict: Optional[bool] = None,
722
+ ) -> Union[Tuple[torch.Tensor], MaskedLMOutput]:
723
+ # labels should be a `torch.LongTensor` of shape
724
+ # `(batch_size, sequence_length)`. These are used for computing the
725
+ # masked language modeling loss.
726
+ #
727
+ # Indices should be in `[-100, 0, ..., config.vocab_size]` (see
728
+ # `input_ids` docstring) Tokens with indices set to `-100` are ignored
729
+ # (masked), the loss is only computed for the tokens with labels in `[0,
730
+ # ..., config.vocab_size]`
731
+ #
732
+ # Prediction scores are only computed for masked tokens and the (bs,
733
+ # seqlen) dimensions are flattened
734
+ if (input_ids is not None) == (inputs_embeds is not None):
735
+ raise ValueError('Must specify either input_ids or input_embeds!')
736
+
737
+ if labels is None:
738
+ masked_tokens_mask = None
739
+ else:
740
+ masked_tokens_mask = labels > 0
741
+
742
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
743
+
744
+ outputs = self.bert(
745
+ input_ids,
746
+ attention_mask=attention_mask,
747
+ token_type_ids=token_type_ids,
748
+ position_ids=position_ids,
749
+ head_mask=head_mask,
750
+ inputs_embeds=inputs_embeds,
751
+ encoder_hidden_states=encoder_hidden_states,
752
+ encoder_attention_mask=encoder_attention_mask,
753
+ output_attentions=output_attentions,
754
+ output_hidden_states=output_hidden_states,
755
+ return_dict=return_dict,
756
+ masked_tokens_mask=masked_tokens_mask,
757
+ )
758
+
759
+ sequence_output = outputs[0]
760
+ prediction_scores = self.cls(sequence_output)
761
+
762
+ loss = None
763
+ if labels is not None:
764
+ # Compute loss
765
+ loss_fct = nn.CrossEntropyLoss()
766
+ masked_token_idx = torch.nonzero(labels.flatten() > 0,
767
+ as_tuple=False).flatten()
768
+ loss = loss_fct(prediction_scores,
769
+ labels.flatten()[masked_token_idx])
770
+
771
+ assert input_ids is not None, 'Coding error; please open an issue'
772
+ batch, seqlen = input_ids.shape[:2]
773
+ prediction_scores = rearrange(index_put_first_axis(
774
+ prediction_scores, masked_token_idx, batch * seqlen),
775
+ '(b s) d -> b s d',
776
+ b=batch)
777
+
778
+ if not return_dict:
779
+ output = (prediction_scores,) + outputs[2:]
780
+ return ((loss,) + output) if loss is not None else output
781
+
782
+ return MaskedLMOutput(
783
+ loss=loss,
784
+ logits=prediction_scores,
785
+ hidden_states=outputs[0],
786
+ attentions=None,
787
+ )
788
+
789
+ def prepare_inputs_for_generation(self, input_ids: torch.Tensor,
790
+ attention_mask: torch.Tensor,
791
+ **model_kwargs):
792
+ input_shape = input_ids.shape
793
+ effective_batch_size = input_shape[0]
794
+
795
+ # add a dummy token
796
+ if self.config.pad_token_id is None:
797
+ raise ValueError('The PAD token should be defined for generation')
798
+
799
+ attention_mask = torch.cat([
800
+ attention_mask,
801
+ attention_mask.new_zeros((attention_mask.shape[0], 1))
802
+ ],
803
+ dim=-1)
804
+ dummy_token = torch.full((effective_batch_size, 1),
805
+ self.config.pad_token_id,
806
+ dtype=torch.long,
807
+ device=input_ids.device)
808
+ input_ids = torch.cat([input_ids, dummy_token], dim=1)
809
+
810
+ return {'input_ids': input_ids, 'attention_mask': attention_mask}
811
+
812
+
813
+ class BertForNextSentencePrediction(BertPreTrainedModel):
814
+ #TBD: Push in future commit
815
+ pass
816
+
817
+
818
+ class BertForSequenceClassification(BertPreTrainedModel):
819
+ """Bert Model transformer with a sequence classification/regression head.
820
+
821
+ This head is just a linear layer on top of the pooled output. Used for,
822
+ e.g., GLUE tasks.
823
+ """
824
+
825
+ def __init__(self, config):
826
+ super().__init__(config)
827
+ self.num_labels = config.num_labels
828
+ self.config = config
829
+
830
+ self.bert = BertModel(config)
831
+ classifier_dropout = (config.classifier_dropout
832
+ if config.classifier_dropout is not None else
833
+ config.hidden_dropout_prob)
834
+ self.dropout = nn.Dropout(classifier_dropout)
835
+ self.classifier = nn.Linear(config.hidden_size, config.num_labels)
836
+
837
+ # Initialize weights and apply final processing
838
+ self.post_init()
839
+
840
+
841
+ def forward(
842
+ self,
843
+ input_ids: Optional[torch.Tensor] = None,
844
+ attention_mask: Optional[torch.Tensor] = None,
845
+ token_type_ids: Optional[torch.Tensor] = None,
846
+ position_ids: Optional[torch.Tensor] = None,
847
+ head_mask: Optional[torch.Tensor] = None,
848
+ inputs_embeds: Optional[torch.Tensor] = None,
849
+ labels: Optional[torch.Tensor] = None,
850
+ output_attentions: Optional[bool] = None,
851
+ output_hidden_states: Optional[bool] = None,
852
+ return_dict: Optional[bool] = None,
853
+ ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:
854
+ # labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
855
+ # Labels for computing the sequence classification/regression loss.
856
+ # Indices should be in `[0, ..., config.num_labels - 1]`.
857
+ # If `config.num_labels == 1` a regression loss is computed
858
+ # (mean-square loss). If `config.num_labels > 1` a classification loss
859
+ # is computed (cross-entropy).
860
+
861
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
862
+
863
+ outputs = self.bert(
864
+ input_ids,
865
+ attention_mask=attention_mask,
866
+ token_type_ids=token_type_ids,
867
+ position_ids=position_ids,
868
+ head_mask=head_mask,
869
+ inputs_embeds=inputs_embeds,
870
+ output_attentions=output_attentions,
871
+ output_hidden_states=output_hidden_states,
872
+ return_dict=return_dict,
873
+ )
874
+
875
+ pooled_output = outputs[1]
876
+
877
+ pooled_output = self.dropout(pooled_output)
878
+ logits = self.classifier(pooled_output)
879
+
880
+ loss = None
881
+ if labels is not None:
882
+ # Compute loss
883
+ if self.config.problem_type is None:
884
+ if self.num_labels == 1:
885
+ self.config.problem_type = 'regression'
886
+ elif self.num_labels > 1 and (labels.dtype == torch.long or
887
+ labels.dtype == torch.int):
888
+ self.config.problem_type = 'single_label_classification'
889
+ else:
890
+ self.config.problem_type = 'multi_label_classification'
891
+
892
+ if self.config.problem_type == 'regression':
893
+ loss_fct = nn.MSELoss()
894
+ if self.num_labels == 1:
895
+ loss = loss_fct(logits.squeeze(), labels.squeeze())
896
+ else:
897
+ loss = loss_fct(logits, labels)
898
+ elif self.config.problem_type == 'single_label_classification':
899
+ loss_fct = nn.CrossEntropyLoss()
900
+ loss = loss_fct(logits.view(-1, self.num_labels),
901
+ labels.view(-1))
902
+ elif self.config.problem_type == 'multi_label_classification':
903
+ loss_fct = nn.BCEWithLogitsLoss()
904
+ loss = loss_fct(logits, labels)
905
+
906
+ if not return_dict:
907
+ output = (logits,) + outputs[2:]
908
+ return ((loss,) + output) if loss is not None else output
909
+
910
+ return SequenceClassifierOutput(
911
+ loss=loss,
912
+ logits=logits,
913
+ hidden_states=outputs[0],
914
+ attentions=None,
915
+ )
916
+
bert_padding.py ADDED
@@ -0,0 +1,154 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2022 MosaicML Examples authors
2
+ # SPDX-License-Identifier: Apache-2.0
3
+
4
+ # Adapted from https://github.com/HazyResearch/flash-attention/blob/main/flash_attn/bert_padding.py
5
+ # Which was adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py
6
+
7
+
8
+ from typing import Tuple, cast
9
+
10
+ import torch
11
+ import torch.nn.functional as F
12
+ from einops import rearrange, repeat
13
+
14
+
15
+ class IndexFirstAxis(torch.autograd.Function):
16
+
17
+ @staticmethod
18
+ def forward(ctx, input: torch.Tensor,
19
+ indices: torch.Tensor) -> torch.Tensor:
20
+ """Get just the values of `input` which are at `indices`.
21
+
22
+ Arguments:
23
+ ctx: the autograd context object
24
+ input: (b, ...) 2+ dimensional tensor
25
+ indices: (num_idx) 1D tensor
26
+ """
27
+ ctx.save_for_backward(indices)
28
+ assert input.ndim >= 2
29
+ ctx.first_axis_dim, other_shape = input.shape[0], input.shape[
30
+ 1:] # type: ignore
31
+ second_dim = other_shape.numel(
32
+ ) # product of sizes of all but first dimension
33
+ # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing.
34
+ return torch.gather(
35
+ rearrange(input, 'b ... -> b (...)'), # (b, ...) -> (b, second_dim)
36
+ 0,
37
+ repeat(indices, 'z -> z d',
38
+ d=second_dim) # (indices,) -> (indices, second_dim)
39
+ ).reshape(-1, *other_shape) # (num_idx, ...)
40
+
41
+ @staticmethod
42
+ def backward(ctx, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None]:
43
+ indices, = ctx.saved_tensors
44
+ assert grad_output.ndim >= 2
45
+ other_shape = grad_output.shape[1:]
46
+ grad_output = rearrange(grad_output, 'b ... -> b (...)')
47
+ grad_input = torch.zeros([ctx.first_axis_dim, grad_output.shape[1]],
48
+ device=grad_output.device,
49
+ dtype=grad_output.dtype)
50
+ # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing.
51
+ # grad_input[indices] = grad_output
52
+ grad_input.scatter_(0,
53
+ repeat(indices, 'z -> z d', d=grad_output.shape[1]),
54
+ grad_output)
55
+ return grad_input.reshape(ctx.first_axis_dim, *other_shape), None
56
+
57
+
58
+ index_first_axis = IndexFirstAxis.apply
59
+
60
+
61
+ class IndexPutFirstAxis(torch.autograd.Function):
62
+
63
+ @staticmethod
64
+ def forward(ctx, values: torch.Tensor, indices: torch.Tensor,
65
+ first_axis_dim) -> torch.Tensor:
66
+ ctx.save_for_backward(indices)
67
+ assert indices.ndim == 1
68
+ assert values.ndim >= 2
69
+ output = torch.zeros(first_axis_dim,
70
+ *values.shape[1:],
71
+ device=values.device,
72
+ dtype=values.dtype)
73
+ output[indices] = values
74
+ return output
75
+
76
+ @staticmethod
77
+ def backward(ctx,
78
+ grad_output: torch.Tensor) -> Tuple[torch.Tensor, None, None]:
79
+ indices, = ctx.saved_tensors
80
+ grad_values = grad_output[indices]
81
+ return grad_values, None, None
82
+
83
+
84
+ index_put_first_axis = IndexPutFirstAxis.apply
85
+
86
+
87
+ def unpad_input(
88
+ hidden_states: torch.Tensor,
89
+ attention_mask: torch.Tensor,
90
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
91
+ """Remove padding from input sequences.
92
+
93
+ Arguments:
94
+ hidden_states: (batch, seqlen, ...)
95
+ attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid.
96
+
97
+ Returns:
98
+ hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
99
+ indices: (total_nnz)
100
+ cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states.
101
+ max_seqlen_in_batch: int ()
102
+ """
103
+ seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
104
+ indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
105
+ max_seqlen_in_batch = int(seqlens_in_batch.max().item())
106
+ cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32),
107
+ (1, 0))
108
+ # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the
109
+ # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim
110
+ # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to
111
+ # index with integer indices. Moreover, torch's index is a bit slower than it needs to be,
112
+ # so we write custom forward and backward to make it a bit faster.
113
+ hidden_states = cast(
114
+ torch.Tensor,
115
+ index_first_axis(rearrange(hidden_states, 'b s ... -> (b s) ...'),
116
+ indices))
117
+ return hidden_states, indices, cu_seqlens, max_seqlen_in_batch
118
+
119
+
120
+ def unpad_input_only(
121
+ hidden_states: torch.Tensor,
122
+ attention_mask: torch.Tensor,
123
+ ) -> torch.Tensor:
124
+ """Like unpad_input, but only return the unpadded first tensor.
125
+
126
+ Save a small amount of overhead.
127
+
128
+ Arguments:
129
+ hidden_states: (batch, seqlen, ...)
130
+ attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid.
131
+
132
+ Returns:
133
+ hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
134
+ """
135
+ indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
136
+ return index_first_axis(rearrange(hidden_states, 'b s ... -> (b s) ...'),
137
+ indices)
138
+
139
+
140
+ def pad_input(hidden_states: torch.Tensor, indices: torch.Tensor, batch: int,
141
+ seqlen: int) -> torch.Tensor:
142
+ """Add padding to sequences.
143
+
144
+ Arguments:
145
+ hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
146
+ indices: (total_nnz)
147
+ batch: int batch_size
148
+ seqlen: int max sequence length
149
+
150
+ Returns:
151
+ hidden_states: (batch, seqlen, ...)
152
+ """
153
+ output = index_put_first_axis(hidden_states, indices, batch * seqlen)
154
+ return rearrange(output, '(b s) ... -> b s ...', b=batch)
config.json ADDED
@@ -0,0 +1,42 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "dnabert2_conservation",
3
+ "alibi_starting_size": 512,
4
+ "architectures": [
5
+ "BertForSequenceClassification"
6
+ ],
7
+ "attention_probs_dropout_prob": 0.0,
8
+ "auto_map": {
9
+ "AutoConfig": "configuration_bert.BertConfig",
10
+ "AutoModel": "bert_layers.BertModel",
11
+ "AutoModelForMaskedLM": "bert_layers.BertForMaskedLM",
12
+ "AutoModelForSequenceClassification": "bert_layers.BertForSequenceClassification"
13
+ },
14
+ "classifier_dropout": null,
15
+ "gradient_checkpointing": false,
16
+ "hidden_act": "gelu",
17
+ "hidden_dropout_prob": 0.1,
18
+ "hidden_size": 768,
19
+ "id2label": {
20
+ "0": "Not conserved",
21
+ "1": "Conserved"
22
+ },
23
+ "initializer_range": 0.02,
24
+ "intermediate_size": 3072,
25
+ "label2id": {
26
+ "Not conserved": 0,
27
+ "Conserved": 1
28
+ },
29
+ "layer_norm_eps": 1e-12,
30
+ "max_position_embeddings": 512,
31
+ "model_type": "bert",
32
+ "num_attention_heads": 12,
33
+ "num_hidden_layers": 12,
34
+ "pad_token_id": 3,
35
+ "position_embedding_type": "absolute",
36
+ "problem_type": "single_label_classification",
37
+ "torch_dtype": "float32",
38
+ "transformers_version": "4.39.1",
39
+ "type_vocab_size": 2,
40
+ "use_cache": true,
41
+ "vocab_size": 4096
42
+ }
configuration_bert.py ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2022 MosaicML Examples authors
2
+ # SPDX-License-Identifier: Apache-2.0
3
+
4
+ from transformers import BertConfig as TransformersBertConfig
5
+
6
+
7
+ class BertConfig(TransformersBertConfig):
8
+
9
+ def __init__(
10
+ self,
11
+ alibi_starting_size: int = 512,
12
+ attention_probs_dropout_prob: float = 0.0,
13
+ **kwargs,
14
+ ):
15
+ """Configuration class for MosaicBert.
16
+
17
+ Args:
18
+ alibi_starting_size (int): Use `alibi_starting_size` to determine how large of an alibi tensor to
19
+ create when initializing the model. You should be able to ignore this parameter in most cases.
20
+ Defaults to 512.
21
+ attention_probs_dropout_prob (float): By default, turn off attention dropout in Mosaic BERT
22
+ (otherwise, Flash Attention will be off by default). Defaults to 0.0.
23
+ """
24
+ super().__init__(
25
+ attention_probs_dropout_prob=attention_probs_dropout_prob, **kwargs)
26
+ self.alibi_starting_size = alibi_starting_size
flash_attn_triton.py ADDED
@@ -0,0 +1,1112 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2022 MosaicML Examples authors
2
+ # SPDX-License-Identifier: Apache-2.0
3
+
4
+ """Triton implementation of Flash Attention.
5
+
6
+ # Copyright (c) 2022, Tri Dao.
7
+ #
8
+ # Licensed under the Apache License, Version 2.0 (the "License");
9
+ # you may not use this file except in compliance with the License.
10
+ # You may obtain a copy of the License at
11
+ #
12
+ # http://www.apache.org/licenses/LICENSE-2.0
13
+ #
14
+ # Unless required by applicable law or agreed to in writing, software
15
+ # distributed under the License is distributed on an "AS IS" BASIS,
16
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17
+ # See the License for the specific language governing permissions and
18
+ # limitations under the License.
19
+
20
+ *Experimental* implementation of FlashAttention in Triton.
21
+ We use the FlashAttention implementation from Phil Tillet a starting point.
22
+ https://github.com/openai/triton/blob/master/python/tutorials/06-fused-attention.py
23
+
24
+ Changes:
25
+ - Implement both causal and non-causal attention.
26
+ - Implement both self-attention and cross-attention.
27
+ - Support arbitrary seqlens (not just multiples of 128), for both forward and backward.
28
+ - Support all head dimensions up to 128 (not just 16, 32, 64, 128), for both forward and backward.
29
+ - Support attention bias.
30
+ - Speed up the forward pass a bit, and only store the LSE instead of m and l.
31
+ - Make the backward for d=128 much faster by reducing register spilling.
32
+ - Optionally parallelize the backward pass across seqlen_k, to deal with the case of
33
+ small batch size * nheads.
34
+
35
+ Caution:
36
+ - If you plan to use headdim other than 64 and 128, you should test for race conditions
37
+ (due to the Triton compiler), as done in tests/test_flash_attn.py
38
+ "test_flash_attn_triton_race_condition". I've tested and fixed many race conditions
39
+ for different head dimensions (40, 48, 64, 128, 80, 88, 96), but I'm still not 100% confident
40
+ that there are none left for other head dimensions.
41
+ Differences between this Triton version and the CUDA version:
42
+ - Triton version doesn't support dropout.
43
+ - Triton forward is generally faster than CUDA forward.
44
+ - Triton backward is faster than CUDA backward when batch * nheads is small, and when headdim=64.
45
+ It is slightly slower when headdim=128 and batch * nheads is large.
46
+ - Triton version doesn't yet support different sequence lengths in a batch (i.e., RaggedTensor/NestedTensor).
47
+ """
48
+
49
+ import math
50
+
51
+ import torch
52
+ import triton # type: ignore (reportMissingImports)
53
+ import triton.language as tl # type: ignore (reportMissingImports)
54
+ from einops import repeat
55
+
56
+
57
+ @triton.autotune(
58
+ configs=[
59
+ triton.Config({
60
+ 'BLOCK_M': 128,
61
+ 'BLOCK_N': 128
62
+ },
63
+ num_warps=8,
64
+ num_stages=1),
65
+ # This config has a race condition when EVEN_M == False, disabling it for now.
66
+ # triton.Config({"BLOCK_M": 64, "BLOCK_N": 64}, num_warps=4, num_stages=1),
67
+ ],
68
+ key=[
69
+ 'CACHE_KEY_SEQLEN_Q', 'CACHE_KEY_SEQLEN_K', 'BIAS_TYPE', 'IS_CAUSAL',
70
+ 'BLOCK_HEADDIM'
71
+ ])
72
+ @triton.heuristics({
73
+ 'EVEN_M': lambda args: args['seqlen_q'] % args['BLOCK_M'] == 0,
74
+ 'EVEN_N': lambda args: args['seqlen_k'] % args['BLOCK_N'] == 0,
75
+ 'EVEN_HEADDIM': lambda args: args['headdim'] == args['BLOCK_HEADDIM'],
76
+ })
77
+ @triton.jit
78
+ def _fwd_kernel(
79
+ Q,
80
+ K,
81
+ V,
82
+ Bias,
83
+ Out,
84
+ Lse,
85
+ TMP, # NOTE: TMP is a scratchpad buffer to workaround a compiler bug
86
+ softmax_scale,
87
+ stride_qb,
88
+ stride_qh,
89
+ stride_qm,
90
+ stride_kb,
91
+ stride_kh,
92
+ stride_kn,
93
+ stride_vb,
94
+ stride_vh,
95
+ stride_vn,
96
+ stride_bb,
97
+ stride_bh,
98
+ stride_bm,
99
+ stride_ob,
100
+ stride_oh,
101
+ stride_om,
102
+ nheads,
103
+ seqlen_q,
104
+ seqlen_k,
105
+ seqlen_q_rounded,
106
+ headdim,
107
+ CACHE_KEY_SEQLEN_Q,
108
+ CACHE_KEY_SEQLEN_K,
109
+ BIAS_TYPE: tl.constexpr,
110
+ IS_CAUSAL: tl.constexpr,
111
+ BLOCK_HEADDIM: tl.constexpr,
112
+ EVEN_M: tl.constexpr,
113
+ EVEN_N: tl.constexpr,
114
+ EVEN_HEADDIM: tl.constexpr,
115
+ BLOCK_M: tl.constexpr,
116
+ BLOCK_N: tl.constexpr,
117
+ ):
118
+ start_m = tl.program_id(0)
119
+ off_hb = tl.program_id(1)
120
+ off_b = off_hb // nheads
121
+ off_h = off_hb % nheads
122
+ # off_b = tl.program_id(1)
123
+ # off_h = tl.program_id(2)
124
+ # off_hb = off_b * nheads + off_h
125
+ # initialize offsets
126
+ offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
127
+ offs_n = tl.arange(0, BLOCK_N)
128
+ offs_d = tl.arange(0, BLOCK_HEADDIM)
129
+ # Initialize pointers to Q, K, V
130
+ # Adding parenthesis around indexing might use int32 math instead of int64 math?
131
+ # https://github.com/openai/triton/issues/741
132
+ # I'm seeing a tiny bit of difference (5-7us)
133
+ q_ptrs = Q + off_b * stride_qb + off_h * stride_qh + (
134
+ offs_m[:, None] * stride_qm + offs_d[None, :])
135
+ k_ptrs = K + off_b * stride_kb + off_h * stride_kh + (
136
+ offs_n[:, None] * stride_kn + offs_d[None, :])
137
+ v_ptrs = V + off_b * stride_vb + off_h * stride_vh + (
138
+ offs_n[:, None] * stride_vn + offs_d[None, :])
139
+ if BIAS_TYPE == 'vector':
140
+ b_ptrs = Bias + off_b * stride_bb + off_h * stride_bh + offs_n
141
+ elif BIAS_TYPE == 'matrix':
142
+ b_ptrs = Bias + off_b * stride_bb + off_h * stride_bh + (
143
+ offs_m[:, None] * stride_bm + offs_n[None, :])
144
+ else:
145
+ raise ValueError("BIAS_TYPE must be one of {'vector', 'matrix'}")
146
+ # initialize pointer to m and l
147
+ t_ptrs = TMP + off_hb * seqlen_q_rounded + offs_m
148
+ lse_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float('inf')
149
+ m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float('inf')
150
+ acc_o = tl.zeros([BLOCK_M, BLOCK_HEADDIM], dtype=tl.float32)
151
+ # load q: it will stay in SRAM throughout
152
+ # [2022-10-30] TD: Triton bug - in the case of EVEN_M=True and EVEN_N=False, if we just call
153
+ # tl.load(q_ptrs), we get the wrong output!
154
+ if EVEN_M & EVEN_N:
155
+ if EVEN_HEADDIM:
156
+ q = tl.load(q_ptrs)
157
+ else:
158
+ q = tl.load(q_ptrs, mask=offs_d[None, :] < headdim, other=0.0)
159
+ else:
160
+ if EVEN_HEADDIM:
161
+ q = tl.load(q_ptrs, mask=offs_m[:, None] < seqlen_q, other=0.0)
162
+ else:
163
+ q = tl.load(q_ptrs,
164
+ mask=(offs_m[:, None] < seqlen_q) &
165
+ (offs_d[None, :] < headdim),
166
+ other=0.0)
167
+ # loop over k, v and update accumulator
168
+ end_n = seqlen_k if not IS_CAUSAL else tl.minimum(
169
+ (start_m + 1) * BLOCK_M, seqlen_k)
170
+ for start_n in range(0, end_n, BLOCK_N):
171
+ start_n = tl.multiple_of(start_n, BLOCK_N)
172
+ # -- compute qk ----
173
+ if EVEN_N & EVEN_M: # If we just do "if EVEN_N", there seems to be some race condition
174
+ if EVEN_HEADDIM:
175
+ k = tl.load(k_ptrs + start_n * stride_kn)
176
+ else:
177
+ k = tl.load(k_ptrs + start_n * stride_kn,
178
+ mask=offs_d[None, :] < headdim,
179
+ other=0.0)
180
+ else:
181
+ if EVEN_HEADDIM:
182
+ k = tl.load(k_ptrs + start_n * stride_kn,
183
+ mask=(start_n + offs_n)[:, None] < seqlen_k,
184
+ other=0.0)
185
+ else:
186
+ k = tl.load(k_ptrs + start_n * stride_kn,
187
+ mask=((start_n + offs_n)[:, None] < seqlen_k) &
188
+ (offs_d[None, :] < headdim),
189
+ other=0.0)
190
+ qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
191
+ qk += tl.dot(q, k, trans_b=True)
192
+ # Trying to combine the two masks seem to make the result wrong
193
+ if not EVEN_N: # Need to mask out otherwise the softmax is wrong
194
+ qk += tl.where((start_n + offs_n)[None, :] < seqlen_k, 0,
195
+ float('-inf'))
196
+ if IS_CAUSAL:
197
+ qk += tl.where(offs_m[:, None] >= (start_n + offs_n)[None, :], 0,
198
+ float('-inf'))
199
+ if BIAS_TYPE != 'none':
200
+ if BIAS_TYPE == 'vector':
201
+ if EVEN_N:
202
+ bias = tl.load(b_ptrs + start_n).to(tl.float32)
203
+ else:
204
+ bias = tl.load(b_ptrs + start_n,
205
+ mask=(start_n + offs_n) < seqlen_k,
206
+ other=0.0).to(tl.float32)
207
+ bias = bias[None, :]
208
+ elif BIAS_TYPE == 'matrix':
209
+ if EVEN_M & EVEN_N:
210
+ bias = tl.load(b_ptrs + start_n).to(tl.float32)
211
+ else:
212
+ bias = tl.load(b_ptrs + start_n,
213
+ mask=(offs_m[:, None] < seqlen_q) &
214
+ ((start_n + offs_n)[None, :] < seqlen_k),
215
+ other=0.0).to(tl.float32)
216
+ else:
217
+ raise ValueError(
218
+ "BIAS_TYPE must be one of {'vector', 'matrix'}")
219
+ # Slightly faster to multiply the softmax_scale in the tl.exp below since the compiler
220
+ # can then fuse the mult and add into an fma instruction. But if we have bias we need to
221
+ # to multiply with softmax_scale here.
222
+ qk = qk * softmax_scale + bias
223
+ m_ij = tl.maximum(tl.max(qk, 1), lse_i)
224
+ p = tl.exp(qk - m_ij[:, None])
225
+ else:
226
+ m_ij = tl.maximum(tl.max(qk, 1) * softmax_scale, lse_i)
227
+ p = tl.exp(qk * softmax_scale - m_ij[:, None])
228
+ l_ij = tl.sum(p, 1)
229
+
230
+ # scale acc_o
231
+ acc_o_scale = tl.exp(m_i - m_ij)
232
+
233
+ # # -- update output accumulator --
234
+ # BUG: have to store and immediately load
235
+ tl.store(t_ptrs, acc_o_scale)
236
+ acc_o_scale = tl.load(t_ptrs)
237
+ acc_o = acc_o * acc_o_scale[:, None]
238
+ # update acc_o
239
+ if EVEN_N & EVEN_M: # If we just do "if EVEN_N", there seems to be some race condition
240
+ if EVEN_HEADDIM:
241
+ v = tl.load(v_ptrs + start_n * stride_vn)
242
+ else:
243
+ v = tl.load(v_ptrs + start_n * stride_vn,
244
+ mask=offs_d[None, :] < headdim,
245
+ other=0.0)
246
+ else:
247
+ if EVEN_HEADDIM:
248
+ v = tl.load(v_ptrs + start_n * stride_vn,
249
+ mask=(start_n + offs_n)[:, None] < seqlen_k,
250
+ other=0.0)
251
+ else:
252
+ v = tl.load(v_ptrs + start_n * stride_vn,
253
+ mask=((start_n + offs_n)[:, None] < seqlen_k) &
254
+ (offs_d[None, :] < headdim),
255
+ other=0.0)
256
+ p = p.to(v.dtype)
257
+ acc_o += tl.dot(p, v)
258
+
259
+ # -- update statistics
260
+ m_i = m_ij
261
+ l_i_new = tl.exp(lse_i - m_ij) + l_ij
262
+ lse_i = m_ij + tl.log(l_i_new)
263
+
264
+ o_scale = tl.exp(m_i - lse_i)
265
+ # BUG: have to store and immediately load
266
+ tl.store(t_ptrs, o_scale)
267
+ o_scale = tl.load(t_ptrs)
268
+ acc_o = acc_o * o_scale[:, None]
269
+ # rematerialize offsets to save registers
270
+ start_m = tl.program_id(0)
271
+ offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
272
+ # write back l and m
273
+ lse_ptrs = Lse + off_hb * seqlen_q_rounded + offs_m
274
+ tl.store(lse_ptrs, lse_i)
275
+ # initialize pointers to output
276
+ offs_n = tl.arange(0, BLOCK_HEADDIM)
277
+ out_ptrs = Out + off_b * stride_ob + off_h * stride_oh + (
278
+ offs_m[:, None] * stride_om + offs_n[None, :])
279
+ if EVEN_M:
280
+ if EVEN_HEADDIM:
281
+ tl.store(out_ptrs, acc_o)
282
+ else:
283
+ tl.store(out_ptrs, acc_o, mask=offs_d[None, :] < headdim)
284
+ else:
285
+ if EVEN_HEADDIM:
286
+ tl.store(out_ptrs, acc_o, mask=offs_m[:, None] < seqlen_q)
287
+ else:
288
+ tl.store(out_ptrs,
289
+ acc_o,
290
+ mask=(offs_m[:, None] < seqlen_q) &
291
+ (offs_d[None, :] < headdim))
292
+
293
+
294
+ @triton.jit
295
+ def _bwd_preprocess_do_o_dot(
296
+ Out,
297
+ DO,
298
+ Delta,
299
+ stride_ob,
300
+ stride_oh,
301
+ stride_om,
302
+ stride_dob,
303
+ stride_doh,
304
+ stride_dom,
305
+ nheads,
306
+ seqlen_q,
307
+ seqlen_q_rounded,
308
+ headdim,
309
+ BLOCK_M: tl.constexpr,
310
+ BLOCK_HEADDIM: tl.constexpr,
311
+ ):
312
+ start_m = tl.program_id(0)
313
+ off_hb = tl.program_id(1)
314
+ off_b = off_hb // nheads
315
+ off_h = off_hb % nheads
316
+ # initialize offsets
317
+ offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
318
+ offs_d = tl.arange(0, BLOCK_HEADDIM)
319
+ # load
320
+ o = tl.load(Out + off_b * stride_ob + off_h * stride_oh +
321
+ offs_m[:, None] * stride_om + offs_d[None, :],
322
+ mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim),
323
+ other=0.0).to(tl.float32)
324
+ do = tl.load(DO + off_b * stride_dob + off_h * stride_doh +
325
+ offs_m[:, None] * stride_dom + offs_d[None, :],
326
+ mask=(offs_m[:, None] < seqlen_q) &
327
+ (offs_d[None, :] < headdim),
328
+ other=0.0).to(tl.float32)
329
+ delta = tl.sum(o * do, axis=1)
330
+ # write-back
331
+ tl.store(Delta + off_hb * seqlen_q_rounded + offs_m, delta)
332
+
333
+
334
+ @triton.jit
335
+ def _bwd_kernel_one_col_block(
336
+ start_n,
337
+ Q,
338
+ K,
339
+ V,
340
+ Bias,
341
+ DO,
342
+ DQ,
343
+ DK,
344
+ DV,
345
+ LSE,
346
+ D,
347
+ softmax_scale,
348
+ stride_qm,
349
+ stride_kn,
350
+ stride_vn,
351
+ stride_bm,
352
+ stride_dom,
353
+ stride_dqm,
354
+ stride_dkn,
355
+ stride_dvn,
356
+ seqlen_q,
357
+ seqlen_k,
358
+ headdim,
359
+ ATOMIC_ADD: tl.constexpr,
360
+ BIAS_TYPE: tl.constexpr,
361
+ IS_CAUSAL: tl.constexpr,
362
+ BLOCK_HEADDIM: tl.constexpr,
363
+ EVEN_M: tl.constexpr,
364
+ EVEN_N: tl.constexpr,
365
+ EVEN_HEADDIM: tl.constexpr,
366
+ BLOCK_M: tl.constexpr,
367
+ BLOCK_N: tl.constexpr,
368
+ ):
369
+ # We need to make sure begin_m is a multiple of BLOCK_M (not BLOCK_N)
370
+ begin_m = 0 if not IS_CAUSAL else ((start_n * BLOCK_N) // BLOCK_M) * BLOCK_M
371
+ # initialize row/col offsets
372
+ offs_qm = begin_m + tl.arange(0, BLOCK_M)
373
+ offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N)
374
+ offs_m = tl.arange(0, BLOCK_M)
375
+ offs_d = tl.arange(0, BLOCK_HEADDIM)
376
+ # initialize pointers to value-like data
377
+ q_ptrs = Q + (offs_qm[:, None] * stride_qm + offs_d[None, :])
378
+ k_ptrs = K + (offs_n[:, None] * stride_kn + offs_d[None, :])
379
+ v_ptrs = V + (offs_n[:, None] * stride_vn + offs_d[None, :])
380
+ do_ptrs = DO + (offs_qm[:, None] * stride_dom + offs_d[None, :])
381
+ dq_ptrs = DQ + (offs_qm[:, None] * stride_dqm + offs_d[None, :])
382
+ if BIAS_TYPE == 'vector':
383
+ b_ptrs = Bias + offs_n
384
+ elif BIAS_TYPE == 'matrix':
385
+ b_ptrs = Bias + (offs_qm[:, None] * stride_bm + offs_n[None, :])
386
+ else:
387
+ raise ValueError("BIAS_TYPE must be one of {'vector', 'matrix'}")
388
+ # initialize dv and dk
389
+ dv = tl.zeros([BLOCK_N, BLOCK_HEADDIM], dtype=tl.float32)
390
+ dk = tl.zeros([BLOCK_N, BLOCK_HEADDIM], dtype=tl.float32)
391
+ # k and v stay in SRAM throughout
392
+ # [2022-10-30] TD: Same bug as the fwd. In the case of EVEN_N=True and EVEN_M=False,
393
+ # if we just call tl.load(k_ptrs), we get the wrong output!
394
+ if EVEN_N & EVEN_M:
395
+ if EVEN_HEADDIM:
396
+ k = tl.load(k_ptrs)
397
+ v = tl.load(v_ptrs)
398
+ else:
399
+ k = tl.load(k_ptrs, mask=offs_d[None, :] < headdim, other=0.0)
400
+ v = tl.load(v_ptrs, mask=offs_d[None, :] < headdim, other=0.0)
401
+ else:
402
+ if EVEN_HEADDIM:
403
+ k = tl.load(k_ptrs, mask=offs_n[:, None] < seqlen_k, other=0.0)
404
+ v = tl.load(v_ptrs, mask=offs_n[:, None] < seqlen_k, other=0.0)
405
+ else:
406
+ k = tl.load(k_ptrs,
407
+ mask=(offs_n[:, None] < seqlen_k) &
408
+ (offs_d[None, :] < headdim),
409
+ other=0.0)
410
+ v = tl.load(v_ptrs,
411
+ mask=(offs_n[:, None] < seqlen_k) &
412
+ (offs_d[None, :] < headdim),
413
+ other=0.0)
414
+ # loop over rows
415
+ num_block_m = tl.cdiv(seqlen_q, BLOCK_M)
416
+ for start_m in range(begin_m, num_block_m * BLOCK_M, BLOCK_M):
417
+ start_m = tl.multiple_of(start_m, BLOCK_M)
418
+ offs_m_curr = start_m + offs_m
419
+ # load q, k, v, do on-chip
420
+ # Same bug as below. Otherwise gives wrong result for headdim=40, seqlen=(128, 117)
421
+ if EVEN_M & EVEN_HEADDIM:
422
+ q = tl.load(q_ptrs)
423
+ else:
424
+ if EVEN_HEADDIM:
425
+ q = tl.load(q_ptrs,
426
+ mask=offs_m_curr[:, None] < seqlen_q,
427
+ other=0.0)
428
+ else:
429
+ q = tl.load(q_ptrs,
430
+ mask=(offs_m_curr[:, None] < seqlen_q) &
431
+ (offs_d[None, :] < headdim),
432
+ other=0.0)
433
+ # recompute p = softmax(qk, dim=-1).T
434
+ qk = tl.dot(q, k, trans_b=True)
435
+ # Trying to combine the two masks seem to make the result wrong
436
+ if not EVEN_N: # Need to mask out otherwise the softmax is wrong
437
+ qk = tl.where(offs_n[None, :] < seqlen_k, qk, float('-inf'))
438
+ if IS_CAUSAL:
439
+ qk = tl.where(offs_m_curr[:, None] >= (offs_n[None, :]), qk,
440
+ float('-inf'))
441
+ if BIAS_TYPE != 'none':
442
+ if BIAS_TYPE == 'vector':
443
+ if EVEN_N:
444
+ bias = tl.load(b_ptrs).to(tl.float32)
445
+ else:
446
+ bias = tl.load(b_ptrs, mask=offs_n < seqlen_k,
447
+ other=0.0).to(tl.float32)
448
+ bias = bias[None, :]
449
+ elif BIAS_TYPE == 'matrix':
450
+ if EVEN_M & EVEN_N:
451
+ bias = tl.load(b_ptrs).to(tl.float32)
452
+ else:
453
+ bias = tl.load(b_ptrs,
454
+ mask=(offs_m_curr[:, None] < seqlen_q) &
455
+ (offs_n[None, :] < seqlen_k),
456
+ other=0.0).to(tl.float32)
457
+ else:
458
+ raise ValueError(
459
+ "BIAS_TYPE must be one of {'vector', 'matrix'}")
460
+ qk = qk * softmax_scale + bias
461
+ # There seems to be a race condition when headdim=48/96, and dq, dk, dv are wrong.
462
+ # Also wrong for headdim=64.
463
+ if not (EVEN_M & EVEN_HEADDIM):
464
+ tl.debug_barrier()
465
+ lse_i = tl.load(LSE + offs_m_curr)
466
+ if BIAS_TYPE == 'none':
467
+ p = tl.exp(qk * softmax_scale - lse_i[:, None])
468
+ else:
469
+ p = tl.exp(qk - lse_i[:, None])
470
+ # compute dv
471
+ # [2022-10-30] TD: A Triton bug: if EVEN_M=True and EVEN_HEADDIM=False, if we call
472
+ # do = tl.load(do_ptrs, mask=offs_d[None, :] < headdim, other=0.0), we get wrong outputs
473
+ # in the case of headdim=48/96, seqlen_q & seqlen_k >= 512. If headdim=40 or seqlen < 512,
474
+ # the output is correct.
475
+ if EVEN_M & EVEN_HEADDIM:
476
+ do = tl.load(do_ptrs)
477
+ else:
478
+ # [2022-11-01] TD: Triton bug, there's a race condition if we just use m_mask and not d_mask.
479
+ do = tl.load(do_ptrs,
480
+ mask=(offs_m_curr[:, None] < seqlen_q) &
481
+ (offs_d[None, :] < headdim),
482
+ other=0.0)
483
+ # if EVEN_M:
484
+ # if EVEN_HEADDIM:
485
+ # do = tl.load(do_ptrs)
486
+ # else:
487
+ # do = tl.load(do_ptrs, mask=offs_d[None, :] < headdim, other=0.0)
488
+ # else:
489
+ # if EVEN_HEADDIM:
490
+ # do = tl.load(do_ptrs, mask=offs_m_curr[:, None] < seqlen_q, other=0.0)
491
+ # else:
492
+ # do = tl.load(do_ptrs, mask=(offs_m_curr[:, None] < seqlen_q)
493
+ # & (offs_d[None, :] < headdim), other=0.0)
494
+ dv += tl.dot(p.to(do.dtype), do, trans_a=True)
495
+ # compute dp = dot(v, do)
496
+ # There seems to be a race condition when headdim=48/96, and dq, dk are wrong.
497
+ # Also wrong for headdim=128, seqlen=(108, 256), and ATOMIC_ADD=True
498
+ # Also wrong for headdim=64, seqlen=(1023, 1024), and ATOMIC_ADD=False
499
+ if not (EVEN_M & EVEN_HEADDIM):
500
+ tl.debug_barrier()
501
+ dp = tl.dot(do, v, trans_b=True)
502
+ # There's a race condition for headdim=48
503
+ if not EVEN_HEADDIM:
504
+ tl.debug_barrier()
505
+ # compute ds = p * (dp - delta[:, None])
506
+ # Putting the subtraction after the dp matmul (instead of before) is slightly faster
507
+ Di = tl.load(D + offs_m_curr)
508
+ # Converting ds to q.dtype here reduces register pressure and makes it much faster
509
+ # for BLOCK_HEADDIM=128
510
+ ds = (p * (dp - Di[:, None]) * softmax_scale).to(q.dtype)
511
+ # compute dk = dot(ds.T, q)
512
+ dk += tl.dot(ds, q, trans_a=True)
513
+ # compute dq
514
+ if not ATOMIC_ADD:
515
+ if EVEN_M & EVEN_HEADDIM: # Race condition if we just do EVEN_M
516
+ dq = tl.load(dq_ptrs, eviction_policy='evict_last')
517
+ dq += tl.dot(ds, k)
518
+ tl.store(dq_ptrs, dq, eviction_policy='evict_last')
519
+ else:
520
+ if EVEN_HEADDIM:
521
+ dq = tl.load(dq_ptrs,
522
+ mask=offs_m_curr[:, None] < seqlen_q,
523
+ other=0.0,
524
+ eviction_policy='evict_last')
525
+ dq += tl.dot(ds, k)
526
+ tl.store(dq_ptrs,
527
+ dq,
528
+ mask=offs_m_curr[:, None] < seqlen_q,
529
+ eviction_policy='evict_last')
530
+ else:
531
+ dq = tl.load(dq_ptrs,
532
+ mask=(offs_m_curr[:, None] < seqlen_q) &
533
+ (offs_d[None, :] < headdim),
534
+ other=0.0,
535
+ eviction_policy='evict_last')
536
+ dq += tl.dot(ds, k)
537
+ tl.store(dq_ptrs,
538
+ dq,
539
+ mask=(offs_m_curr[:, None] < seqlen_q) &
540
+ (offs_d[None, :] < headdim),
541
+ eviction_policy='evict_last')
542
+ else: # If we're parallelizing across the seqlen_k dimension
543
+ dq = tl.dot(ds, k)
544
+ if EVEN_M & EVEN_HEADDIM: # Race condition if we just do EVEN_M
545
+ tl.atomic_add(dq_ptrs, dq)
546
+ else:
547
+ if EVEN_HEADDIM:
548
+ tl.atomic_add(dq_ptrs,
549
+ dq,
550
+ mask=offs_m_curr[:, None] < seqlen_q)
551
+ else:
552
+ tl.atomic_add(dq_ptrs,
553
+ dq,
554
+ mask=(offs_m_curr[:, None] < seqlen_q) &
555
+ (offs_d[None, :] < headdim))
556
+ # increment pointers
557
+ dq_ptrs += BLOCK_M * stride_dqm
558
+ q_ptrs += BLOCK_M * stride_qm
559
+ do_ptrs += BLOCK_M * stride_dom
560
+ if BIAS_TYPE == 'matrix':
561
+ b_ptrs += BLOCK_M * stride_bm
562
+ # write-back
563
+ dv_ptrs = DV + (offs_n[:, None] * stride_dvn + offs_d[None, :])
564
+ dk_ptrs = DK + (offs_n[:, None] * stride_dkn + offs_d[None, :])
565
+ # [2022-11-01] TD: Same bug. In the case of EVEN_N=True and EVEN_M=False,
566
+ # if we just call tl.store(dv_ptrs), there's a race condition
567
+ if EVEN_N & EVEN_M:
568
+ if EVEN_HEADDIM:
569
+ tl.store(dv_ptrs, dv)
570
+ tl.store(dk_ptrs, dk)
571
+ else:
572
+ tl.store(dv_ptrs, dv, mask=offs_d[None, :] < headdim)
573
+ tl.store(dk_ptrs, dk, mask=offs_d[None, :] < headdim)
574
+ else:
575
+ if EVEN_HEADDIM:
576
+ tl.store(dv_ptrs, dv, mask=offs_n[:, None] < seqlen_k)
577
+ tl.store(dk_ptrs, dk, mask=offs_n[:, None] < seqlen_k)
578
+ else:
579
+ tl.store(dv_ptrs,
580
+ dv,
581
+ mask=(offs_n[:, None] < seqlen_k) &
582
+ (offs_d[None, :] < headdim))
583
+ tl.store(dk_ptrs,
584
+ dk,
585
+ mask=(offs_n[:, None] < seqlen_k) &
586
+ (offs_d[None, :] < headdim))
587
+
588
+
589
+ def init_to_zero(name):
590
+ return lambda nargs: nargs[name].zero_()
591
+
592
+
593
+ @triton.autotune(
594
+ configs=[
595
+ triton.Config(
596
+ {
597
+ 'BLOCK_M': 128,
598
+ 'BLOCK_N': 128,
599
+ 'SEQUENCE_PARALLEL': False
600
+ },
601
+ num_warps=8,
602
+ num_stages=1,
603
+ pre_hook=init_to_zero('DQ')),
604
+ triton.Config(
605
+ {
606
+ 'BLOCK_M': 128,
607
+ 'BLOCK_N': 128,
608
+ 'SEQUENCE_PARALLEL': True
609
+ },
610
+ num_warps=8,
611
+ num_stages=1,
612
+ pre_hook=init_to_zero('DQ')),
613
+ # Other configs seem to give wrong results when seqlen_q % 128 != 0, disabling them for now
614
+ # # Kernel is buggy (give wrong result) if we set BLOCK_m=128, BLOCK_n=64, num_warps=*4*
615
+ # triton.Config({"BLOCK_M": 128, "BLOCK_N": 64, "SEQUENCE_PARALLEL": False}, num_warps=8, num_stages=1, pre_hook=init_to_zero('DQ')),
616
+ # triton.Config({"BLOCK_M": 128, "BLOCK_N": 64, "SEQUENCE_PARALLEL": True}, num_warps=8, num_stages=1, pre_hook=init_to_zero('DQ')),
617
+ # triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "SEQUENCE_PARALLEL": False}, num_warps=4, num_stages=1, pre_hook=init_to_zero('DQ')),
618
+ # triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "SEQUENCE_PARALLEL": True}, num_warps=4, num_stages=1, pre_hook=init_to_zero('DQ')),
619
+ ],
620
+ key=[
621
+ 'CACHE_KEY_SEQLEN_Q', 'CACHE_KEY_SEQLEN_K', 'BIAS_TYPE', 'IS_CAUSAL',
622
+ 'BLOCK_HEADDIM'
623
+ ],
624
+ )
625
+ @triton.heuristics({
626
+ 'EVEN_M': lambda args: args['seqlen_q'] % args['BLOCK_M'] == 0,
627
+ 'EVEN_N': lambda args: args['seqlen_k'] % args['BLOCK_N'] == 0,
628
+ 'EVEN_HEADDIM': lambda args: args['headdim'] == args['BLOCK_HEADDIM'],
629
+ })
630
+ @triton.jit
631
+ def _bwd_kernel(
632
+ Q,
633
+ K,
634
+ V,
635
+ Bias,
636
+ DO,
637
+ DQ,
638
+ DK,
639
+ DV,
640
+ LSE,
641
+ D,
642
+ softmax_scale,
643
+ stride_qb,
644
+ stride_qh,
645
+ stride_qm,
646
+ stride_kb,
647
+ stride_kh,
648
+ stride_kn,
649
+ stride_vb,
650
+ stride_vh,
651
+ stride_vn,
652
+ stride_bb,
653
+ stride_bh,
654
+ stride_bm,
655
+ stride_dob,
656
+ stride_doh,
657
+ stride_dom,
658
+ stride_dqb,
659
+ stride_dqh,
660
+ stride_dqm,
661
+ stride_dkb,
662
+ stride_dkh,
663
+ stride_dkn,
664
+ stride_dvb,
665
+ stride_dvh,
666
+ stride_dvn,
667
+ nheads,
668
+ seqlen_q,
669
+ seqlen_k,
670
+ seqlen_q_rounded,
671
+ headdim,
672
+ CACHE_KEY_SEQLEN_Q,
673
+ CACHE_KEY_SEQLEN_K,
674
+ BIAS_TYPE: tl.constexpr,
675
+ IS_CAUSAL: tl.constexpr,
676
+ BLOCK_HEADDIM: tl.constexpr,
677
+ SEQUENCE_PARALLEL: tl.constexpr,
678
+ EVEN_M: tl.constexpr,
679
+ EVEN_N: tl.constexpr,
680
+ EVEN_HEADDIM: tl.constexpr,
681
+ BLOCK_M: tl.constexpr,
682
+ BLOCK_N: tl.constexpr,
683
+ ):
684
+ off_hb = tl.program_id(1)
685
+ off_b = off_hb // nheads
686
+ off_h = off_hb % nheads
687
+ # offset pointers for batch/head
688
+ Q += off_b * stride_qb + off_h * stride_qh
689
+ K += off_b * stride_kb + off_h * stride_kh
690
+ V += off_b * stride_vb + off_h * stride_vh
691
+ DO += off_b * stride_dob + off_h * stride_doh
692
+ DQ += off_b * stride_dqb + off_h * stride_dqh
693
+ DK += off_b * stride_dkb + off_h * stride_dkh
694
+ DV += off_b * stride_dvb + off_h * stride_dvh
695
+ if BIAS_TYPE != 'none':
696
+ Bias += off_b * stride_bb + off_h * stride_bh
697
+ # pointer to row-wise quantities in value-like data
698
+ D += off_hb * seqlen_q_rounded
699
+ LSE += off_hb * seqlen_q_rounded
700
+ if not SEQUENCE_PARALLEL:
701
+ num_block_n = tl.cdiv(seqlen_k, BLOCK_N)
702
+ for start_n in range(0, num_block_n):
703
+ _bwd_kernel_one_col_block(start_n,
704
+ Q,
705
+ K,
706
+ V,
707
+ Bias,
708
+ DO,
709
+ DQ,
710
+ DK,
711
+ DV,
712
+ LSE,
713
+ D,
714
+ softmax_scale,
715
+ stride_qm,
716
+ stride_kn,
717
+ stride_vn,
718
+ stride_bm,
719
+ stride_dom,
720
+ stride_dqm,
721
+ stride_dkn,
722
+ stride_dvn,
723
+ seqlen_q,
724
+ seqlen_k,
725
+ headdim,
726
+ ATOMIC_ADD=False,
727
+ BIAS_TYPE=BIAS_TYPE,
728
+ IS_CAUSAL=IS_CAUSAL,
729
+ BLOCK_HEADDIM=BLOCK_HEADDIM,
730
+ EVEN_M=EVEN_M,
731
+ EVEN_N=EVEN_N,
732
+ EVEN_HEADDIM=EVEN_HEADDIM,
733
+ BLOCK_M=BLOCK_M,
734
+ BLOCK_N=BLOCK_N)
735
+ else:
736
+ start_n = tl.program_id(0)
737
+ _bwd_kernel_one_col_block(start_n,
738
+ Q,
739
+ K,
740
+ V,
741
+ Bias,
742
+ DO,
743
+ DQ,
744
+ DK,
745
+ DV,
746
+ LSE,
747
+ D,
748
+ softmax_scale,
749
+ stride_qm,
750
+ stride_kn,
751
+ stride_vn,
752
+ stride_bm,
753
+ stride_dom,
754
+ stride_dqm,
755
+ stride_dkn,
756
+ stride_dvn,
757
+ seqlen_q,
758
+ seqlen_k,
759
+ headdim,
760
+ ATOMIC_ADD=True,
761
+ BIAS_TYPE=BIAS_TYPE,
762
+ IS_CAUSAL=IS_CAUSAL,
763
+ BLOCK_HEADDIM=BLOCK_HEADDIM,
764
+ EVEN_M=EVEN_M,
765
+ EVEN_N=EVEN_N,
766
+ EVEN_HEADDIM=EVEN_HEADDIM,
767
+ BLOCK_M=BLOCK_M,
768
+ BLOCK_N=BLOCK_N)
769
+
770
+
771
+ def _flash_attn_forward(q, k, v, bias=None, causal=False, softmax_scale=None):
772
+ # shape constraints
773
+ batch, seqlen_q, nheads, d = q.shape
774
+ _, seqlen_k, _, _ = k.shape
775
+ assert k.shape == (batch, seqlen_k, nheads, d)
776
+ assert v.shape == (batch, seqlen_k, nheads, d)
777
+ assert d <= 128, 'FlashAttention only support head dimensions up to 128'
778
+ assert q.dtype == k.dtype == v.dtype, 'All tensors must have the same type'
779
+ assert q.dtype in [torch.float16,
780
+ torch.bfloat16], 'Only support fp16 and bf16'
781
+ assert q.is_cuda and k.is_cuda and v.is_cuda
782
+ softmax_scale = softmax_scale or 1.0 / math.sqrt(d)
783
+
784
+ has_bias = bias is not None
785
+ bias_type = 'none'
786
+ if has_bias:
787
+ assert bias.dtype in [q.dtype, torch.float]
788
+ assert bias.is_cuda
789
+ assert bias.dim() == 4
790
+ if bias.stride(-1) != 1:
791
+ bias = bias.contiguous()
792
+ if bias.shape[2:] == (1, seqlen_k):
793
+ bias_type = 'vector'
794
+ elif bias.shape[2:] == (seqlen_q, seqlen_k):
795
+ bias_type = 'matrix'
796
+ else:
797
+ raise RuntimeError('Last 2 dimensions of bias must be (1, seqlen_k)'
798
+ ' or (seqlen_q, seqlen_k)')
799
+ if bias.shape[:2] == (1, nheads):
800
+ bias = repeat(bias, '1 h ... -> b h ...', b=batch)
801
+ elif bias.shape[:2] == (batch, 1):
802
+ bias = repeat(bias, 'b 1 ... -> b h ...', h=nheads)
803
+ elif bias.shape[:2] == (1, 1):
804
+ bias = repeat(bias, '1 h ... -> b h ...', b=batch)
805
+ bias = repeat(bias, 'b 1 ... -> b h ...', h=nheads)
806
+ assert bias.shape[:2] == (
807
+ batch, nheads
808
+ ), f'First 2 dimensions of bias must be broadcastible to (batch, nheads) = ({batch, nheads}). Bias has shape: {bias.shape}'
809
+ assert bias is not None # for type checking
810
+ bias_strides = (bias.stride(0), bias.stride(1),
811
+ bias.stride(2)) if has_bias else (0, 0, 0)
812
+
813
+ seqlen_q_rounded = math.ceil(seqlen_q / 128) * 128
814
+ lse = torch.empty((batch, nheads, seqlen_q_rounded),
815
+ device=q.device,
816
+ dtype=torch.float32)
817
+ tmp = torch.empty((batch, nheads, seqlen_q_rounded),
818
+ device=q.device,
819
+ dtype=torch.float32)
820
+ o = torch.empty_like(q)
821
+
822
+ BLOCK_HEADDIM = max(triton.next_power_of_2(d), 16)
823
+ # BLOCK = 128
824
+ # num_warps = 4 if d <= 64 else 8
825
+ grid = lambda META: (triton.cdiv(seqlen_q, META['BLOCK_M']), batch * nheads)
826
+ _fwd_kernel[grid]( # type: ignore
827
+ q,
828
+ k,
829
+ v,
830
+ bias,
831
+ o,
832
+ lse,
833
+ tmp,
834
+ softmax_scale,
835
+ q.stride(0),
836
+ q.stride(2),
837
+ q.stride(1),
838
+ k.stride(0),
839
+ k.stride(2),
840
+ k.stride(1),
841
+ v.stride(0),
842
+ v.stride(2),
843
+ v.stride(1),
844
+ *bias_strides,
845
+ o.stride(0),
846
+ o.stride(2),
847
+ o.stride(1),
848
+ nheads,
849
+ seqlen_q,
850
+ seqlen_k,
851
+ seqlen_q_rounded,
852
+ d,
853
+ seqlen_q // 32,
854
+ seqlen_k // 32, # key for triton cache (limit number of compilations)
855
+ # Can't use kwargs here because triton autotune expects key to be args, not kwargs
856
+ # IS_CAUSAL=causal, BLOCK_HEADDIM=d,
857
+ bias_type,
858
+ causal,
859
+ BLOCK_HEADDIM,
860
+ # BLOCK_M=BLOCK, BLOCK_N=BLOCK,
861
+ # num_warps=num_warps,
862
+ # num_stages=1,
863
+ )
864
+ return o, lse, softmax_scale # softmax_scale could have been updated
865
+
866
+
867
+ def _flash_attn_backward(do,
868
+ q,
869
+ k,
870
+ v,
871
+ o,
872
+ lse,
873
+ dq,
874
+ dk,
875
+ dv,
876
+ bias=None,
877
+ causal=False,
878
+ softmax_scale=None):
879
+ # Make sure that the last dimension is contiguous
880
+ if do.stride(-1) != 1:
881
+ do = do.contiguous()
882
+ batch, seqlen_q, nheads, d = q.shape
883
+ _, seqlen_k, _, _ = k.shape
884
+ # assert d in {16, 32, 64, 128}
885
+ assert d <= 128
886
+ seqlen_q_rounded = math.ceil(seqlen_q / 128) * 128
887
+ assert lse.shape == (batch, nheads, seqlen_q_rounded)
888
+ assert q.stride(-1) == k.stride(-1) == v.stride(-1) == o.stride(-1) == 1
889
+ assert dq.stride(-1) == dk.stride(-1) == dv.stride(-1) == 1
890
+ softmax_scale = softmax_scale or 1.0 / math.sqrt(d)
891
+ # dq_accum = torch.zeros_like(q, dtype=torch.float32)
892
+ dq_accum = torch.empty_like(q, dtype=torch.float32)
893
+ delta = torch.empty_like(lse)
894
+ # delta = torch.zeros_like(lse)
895
+
896
+ BLOCK_HEADDIM = max(triton.next_power_of_2(d), 16)
897
+ grid = lambda META: (triton.cdiv(seqlen_q, META['BLOCK_M']), batch * nheads)
898
+ _bwd_preprocess_do_o_dot[grid]( # type: ignore
899
+ o,
900
+ do,
901
+ delta,
902
+ o.stride(0),
903
+ o.stride(2),
904
+ o.stride(1),
905
+ do.stride(0),
906
+ do.stride(2),
907
+ do.stride(1),
908
+ nheads,
909
+ seqlen_q,
910
+ seqlen_q_rounded,
911
+ d,
912
+ BLOCK_M=128,
913
+ BLOCK_HEADDIM=BLOCK_HEADDIM,
914
+ )
915
+
916
+ has_bias = bias is not None
917
+ bias_type = 'none'
918
+ if has_bias:
919
+ assert bias.dtype in [q.dtype, torch.float]
920
+ assert bias.is_cuda
921
+ assert bias.dim() == 4
922
+ assert bias.stride(-1) == 1
923
+ if bias.shape[2:] == (1, seqlen_k):
924
+ bias_type = 'vector'
925
+ elif bias.shape[2:] == (seqlen_q, seqlen_k):
926
+ bias_type = 'matrix'
927
+ else:
928
+ raise RuntimeError('Last 2 dimensions of bias must be (1, seqlen_k)'
929
+ ' or (seqlen_q, seqlen_k)')
930
+ if bias.shape[:2] == (1, nheads):
931
+ bias = repeat(bias, '1 h ... -> b h ...', b=batch)
932
+ elif bias.shape[:2] == (batch, 1):
933
+ bias = repeat(bias, 'b 1 ... -> b h ...', h=nheads)
934
+ elif bias.shape[:2] == (1, 1):
935
+ bias = repeat(bias, '1 h ... -> b h ...', b=batch)
936
+ bias = repeat(bias, 'b 1 ... -> b h ...', h=nheads)
937
+ assert bias.shape[:2] == (
938
+ batch, nheads
939
+ ), f'First 2 dimensions of bias must be broadcastible to (batch, nheads) = ({batch, nheads}). Bias has shape: {bias.shape}'
940
+ assert bias is not None # type checking
941
+ bias_strides = (bias.stride(0), bias.stride(1),
942
+ bias.stride(2)) if has_bias else (0, 0, 0)
943
+
944
+ # BLOCK_M = 128
945
+ # BLOCK_N = 64
946
+ # num_warps = 4
947
+ grid = lambda META: (triton.cdiv(seqlen_k, META['BLOCK_N'])
948
+ if META['SEQUENCE_PARALLEL'] else 1, batch * nheads)
949
+ _bwd_kernel[grid]( # type: ignore
950
+ q,
951
+ k,
952
+ v,
953
+ bias,
954
+ do,
955
+ dq_accum,
956
+ dk,
957
+ dv,
958
+ lse,
959
+ delta,
960
+ softmax_scale,
961
+ q.stride(0),
962
+ q.stride(2),
963
+ q.stride(1),
964
+ k.stride(0),
965
+ k.stride(2),
966
+ k.stride(1),
967
+ v.stride(0),
968
+ v.stride(2),
969
+ v.stride(1),
970
+ *bias_strides,
971
+ do.stride(0),
972
+ do.stride(2),
973
+ do.stride(1),
974
+ dq_accum.stride(0),
975
+ dq_accum.stride(2),
976
+ dq_accum.stride(1),
977
+ dk.stride(0),
978
+ dk.stride(2),
979
+ dk.stride(1),
980
+ dv.stride(0),
981
+ dv.stride(2),
982
+ dv.stride(1),
983
+ nheads,
984
+ seqlen_q,
985
+ seqlen_k,
986
+ seqlen_q_rounded,
987
+ d,
988
+ seqlen_q // 32,
989
+ seqlen_k // 32, # key for triton cache (limit number of compilations)
990
+ # Can't use kwargs here because triton autotune expects key to be args, not kwargs
991
+ # IS_CAUSAL=causal, BLOCK_HEADDIM=d,
992
+ bias_type,
993
+ causal,
994
+ BLOCK_HEADDIM,
995
+ # SEQUENCE_PARALLEL=False,
996
+ # BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N,
997
+ # num_warps=num_warps,
998
+ # num_stages=1,
999
+ )
1000
+ dq.copy_(dq_accum)
1001
+
1002
+
1003
+ class _FlashAttnQKVPackedFunc(torch.autograd.Function):
1004
+
1005
+ @staticmethod
1006
+ def forward(ctx, qkv, bias=None, causal=False, softmax_scale=None):
1007
+ """Forward pass for packed FlashAttention.
1008
+
1009
+ Args:
1010
+ ctx: autograd context
1011
+ qkv: (batch, seqlen, 3, nheads, headdim)
1012
+ bias: optional, shape broadcastible to (batch, nheads, seqlen, seqlen).
1013
+ For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen).
1014
+ ALiBi mask for non-causal would have shape (1, nheads, seqlen, seqlen)
1015
+ causal (bool): whether to incorporate causal attention masking
1016
+ softmax_scale (float, optional): scale factor for softmax
1017
+ """
1018
+ # Make sure that the last dimension is contiguous
1019
+ if qkv.stride(-1) != 1:
1020
+ qkv = qkv.contiguous()
1021
+ o, lse, ctx.softmax_scale = _flash_attn_forward(
1022
+ qkv[:, :, 0],
1023
+ qkv[:, :, 1],
1024
+ qkv[:, :, 2],
1025
+ bias=bias,
1026
+ causal=causal,
1027
+ softmax_scale=softmax_scale)
1028
+ ctx.save_for_backward(qkv, o, lse, bias)
1029
+ ctx.causal = causal
1030
+ return o
1031
+
1032
+ @staticmethod
1033
+ def backward(ctx, do):
1034
+ qkv, o, lse, bias = ctx.saved_tensors
1035
+ assert not ctx.needs_input_grad[
1036
+ 1], 'FlashAttention does not support bias gradient yet'
1037
+ # Triton's autotune causes the Tensor._version to change, and so Pytorch autograd
1038
+ # does a memcpy. To avoid this we run in inference_mode, which doesn't track the version.
1039
+ with torch.inference_mode():
1040
+ dqkv = torch.empty_like(qkv)
1041
+ _flash_attn_backward(do,
1042
+ qkv[:, :, 0],
1043
+ qkv[:, :, 1],
1044
+ qkv[:, :, 2],
1045
+ o,
1046
+ lse,
1047
+ dqkv[:, :, 0],
1048
+ dqkv[:, :, 1],
1049
+ dqkv[:, :, 2],
1050
+ bias=bias,
1051
+ causal=ctx.causal,
1052
+ softmax_scale=ctx.softmax_scale)
1053
+ return dqkv, None, None, None
1054
+
1055
+
1056
+ flash_attn_qkvpacked_func = _FlashAttnQKVPackedFunc.apply
1057
+
1058
+
1059
+ class _FlashAttnFunc(torch.autograd.Function):
1060
+
1061
+ @staticmethod
1062
+ def forward(ctx, q, k, v, bias=None, causal=False, softmax_scale=None):
1063
+ """Forward pass for FlashAttention.
1064
+
1065
+ Args:
1066
+ ctx: autograd context
1067
+ q: (batch_size, seqlen_q, nheads, headdim)
1068
+ k: (batch_size, seqlen_k, nheads, headdim)
1069
+ v: (batch_size, seqlen_k, nheads, headdim)
1070
+ bias: optional, shape broadcastible to (batch, nheads, seqlen_q, seqlen_k).
1071
+ For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen_k).
1072
+ ALiBi mask for non-causal would have shape (1, nheads, seqlen_q, seqlen_k)
1073
+ causal (bool): whether to incorporate causal attention masking
1074
+ softmax_scale (float, optional): scale factor for softmax
1075
+ """
1076
+ # Make sure that the last dimension is contiguous
1077
+ q, k, v = [
1078
+ x if x.stride(-1) == 1 else x.contiguous() for x in [q, k, v]
1079
+ ]
1080
+ o, lse, ctx.softmax_scale = _flash_attn_forward(
1081
+ q, k, v, bias=bias, causal=causal, softmax_scale=softmax_scale)
1082
+ ctx.save_for_backward(q, k, v, o, lse, bias)
1083
+ ctx.causal = causal
1084
+ return o
1085
+
1086
+ @staticmethod
1087
+ def backward(ctx, do):
1088
+ q, k, v, o, lse, bias = ctx.saved_tensors
1089
+ assert not ctx.needs_input_grad[
1090
+ 3], 'FlashAttention does not support bias gradient yet'
1091
+ # Triton's autotune causes the Tensor._version to change, and so Pytorch autograd
1092
+ # does a memcpy. To avoid this we run in inference_mode, which doesn't track the version.
1093
+ with torch.inference_mode():
1094
+ dq = torch.empty_like(q)
1095
+ dk = torch.empty_like(k)
1096
+ dv = torch.empty_like(v)
1097
+ _flash_attn_backward(do,
1098
+ q,
1099
+ k,
1100
+ v,
1101
+ o,
1102
+ lse,
1103
+ dq,
1104
+ dk,
1105
+ dv,
1106
+ bias=bias,
1107
+ causal=ctx.causal,
1108
+ softmax_scale=ctx.softmax_scale)
1109
+ return dq, dk, dv, None, None, None
1110
+
1111
+
1112
+ flash_attn_func = _FlashAttnFunc.apply
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:dad1db886f64247a9e3c4e998c052da028b9994adc9b4b813ac397432bdbb573
3
+ size 468296408
special_tokens_map.json ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "cls_token": {
3
+ "content": "[CLS]",
4
+ "lstrip": false,
5
+ "normalized": false,
6
+ "rstrip": false,
7
+ "single_word": false
8
+ },
9
+ "mask_token": {
10
+ "content": "[MASK]",
11
+ "lstrip": false,
12
+ "normalized": false,
13
+ "rstrip": false,
14
+ "single_word": false
15
+ },
16
+ "pad_token": {
17
+ "content": "[PAD]",
18
+ "lstrip": false,
19
+ "normalized": false,
20
+ "rstrip": false,
21
+ "single_word": false
22
+ },
23
+ "sep_token": {
24
+ "content": "[SEP]",
25
+ "lstrip": false,
26
+ "normalized": false,
27
+ "rstrip": false,
28
+ "single_word": false
29
+ },
30
+ "unk_token": {
31
+ "content": "[UNK]",
32
+ "lstrip": false,
33
+ "normalized": false,
34
+ "rstrip": false,
35
+ "single_word": false
36
+ }
37
+ }
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer_config.json ADDED
@@ -0,0 +1,52 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "added_tokens_decoder": {
3
+ "0": {
4
+ "content": "[UNK]",
5
+ "lstrip": false,
6
+ "normalized": false,
7
+ "rstrip": false,
8
+ "single_word": false,
9
+ "special": true
10
+ },
11
+ "1": {
12
+ "content": "[CLS]",
13
+ "lstrip": false,
14
+ "normalized": false,
15
+ "rstrip": false,
16
+ "single_word": false,
17
+ "special": true
18
+ },
19
+ "2": {
20
+ "content": "[SEP]",
21
+ "lstrip": false,
22
+ "normalized": false,
23
+ "rstrip": false,
24
+ "single_word": false,
25
+ "special": true
26
+ },
27
+ "3": {
28
+ "content": "[PAD]",
29
+ "lstrip": false,
30
+ "normalized": false,
31
+ "rstrip": false,
32
+ "single_word": false,
33
+ "special": true
34
+ },
35
+ "4": {
36
+ "content": "[MASK]",
37
+ "lstrip": false,
38
+ "normalized": false,
39
+ "rstrip": false,
40
+ "single_word": false,
41
+ "special": true
42
+ }
43
+ },
44
+ "clean_up_tokenization_spaces": true,
45
+ "cls_token": "[CLS]",
46
+ "mask_token": "[MASK]",
47
+ "model_max_length": 512,
48
+ "pad_token": "[PAD]",
49
+ "sep_token": "[SEP]",
50
+ "tokenizer_class": "PreTrainedTokenizerFast",
51
+ "unk_token": "[UNK]"
52
+ }