jacobfulano commited on
Commit
51f2ae5
1 Parent(s): 779c4a9

Upload BertForMaskedLM

Browse files
Files changed (6) hide show
  1. bert_layers.py +1072 -0
  2. bert_padding.py +159 -0
  3. config.json +31 -0
  4. configuration_bert.py +25 -0
  5. flash_attn_triton.py +1112 -0
  6. pytorch_model.bin +3 -0
bert_layers.py ADDED
@@ -0,0 +1,1072 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ """Implements Mosaic BERT, with an eye towards the Hugging Face API.
9
+
10
+ Mosaic BERT improves performance over Hugging Face BERT through the following:
11
+
12
+ 1. ALiBi. This architectural change removes positional embeddings and instead encodes positional
13
+ information through attention biases based on query-key position distance. It improves the effectiveness
14
+ of training with shorter sequence lengths by enabling extrapolation to longer sequences.
15
+
16
+ 2. Gated Linear Units (GLU). This architectural change replaces the FFN component of the BERT layer
17
+ to improve overall expressiveness, providing better convergence properties.
18
+
19
+ 3. Flash Attention. The Mosaic BERT's self-attention layer makes use of Flash Attention, which dramatically
20
+ improves the speed of self-attention. Our implementation utilizes a bleeding edge implementation that
21
+ supports attention biases, which allows us to use Flash Attention with ALiBi.
22
+
23
+ 4. Unpadding. Padding is often used to simplify batching across sequences of different lengths. Standard BERT
24
+ implementations waste computation on padded tokens. Mosaic BERT internally unpads to reduce unnecessary computation
25
+ and improve speed. It does this without changing how the user interfaces with the model, thereby
26
+ preserving the simple API of standard implementations.
27
+
28
+
29
+ Currently, Mosaic BERT is available for masked language modeling :class:`BertForMaskedLM` and sequence
30
+ classification :class:`BertForSequenceClassification`. We aim to expand this catalogue in future releases.
31
+
32
+ See :file:`./mosaic_bert.py` for utilities to simplify working with Mosaic BERT in Composer, and for example usage
33
+ of the core Mosaic BERT classes.
34
+ """
35
+
36
+ import copy
37
+ import logging
38
+ import math
39
+ import warnings
40
+ from typing import List, Optional, Tuple, Union
41
+
42
+ import torch
43
+ import torch.nn as nn
44
+ from einops import rearrange
45
+ from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present
46
+ from transformers.activations import ACT2FN
47
+ from transformers.modeling_outputs import (MaskedLMOutput,
48
+ SequenceClassifierOutput)
49
+ from transformers.models.bert.modeling_bert import BertPreTrainedModel
50
+
51
+ from .bert_padding import (index_first_axis,
52
+ index_put_first_axis, pad_input,
53
+ unpad_input, unpad_input_only)
54
+
55
+ try:
56
+ from .flash_attn_triton import flash_attn_qkvpacked_func
57
+ except ImportError as e:
58
+ flash_attn_qkvpacked_func = None
59
+
60
+ logger = logging.getLogger(__name__)
61
+
62
+
63
+ class BertEmbeddings(nn.Module):
64
+ """Construct the embeddings for words, ignoring position.
65
+
66
+ There are no positional embeddings since we use ALiBi and token_type
67
+ embeddings.
68
+
69
+ This module is modeled after the Hugging Face BERT's
70
+ :class:`~transformers.model.bert.modeling_bert.BertEmbeddings`, but is
71
+ modified as part of Mosaic BERT's ALiBi implementation. The key change is
72
+ that position embeddings are removed. Position information instead comes
73
+ from attention biases that scale linearly with the position distance
74
+ between query and key tokens.
75
+
76
+ This module ignores the `position_ids` input to the `forward` method.
77
+ """
78
+
79
+ def __init__(self, config):
80
+ super().__init__()
81
+ self.word_embeddings = nn.Embedding(config.vocab_size,
82
+ config.hidden_size,
83
+ padding_idx=config.pad_token_id)
84
+ # ALiBi doesn't use position embeddings
85
+ self.token_type_embeddings = nn.Embedding(config.type_vocab_size,
86
+ config.hidden_size)
87
+
88
+ # self.LayerNorm is not snake-cased to stick with TensorFlow model
89
+ # variable name and be able to load any TensorFlow checkpoint file
90
+ self.LayerNorm = nn.LayerNorm(config.hidden_size,
91
+ eps=config.layer_norm_eps)
92
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
93
+ self.register_buffer('token_type_ids',
94
+ torch.zeros(config.max_position_embeddings,
95
+ dtype=torch.long),
96
+ persistent=False)
97
+
98
+ def forward(
99
+ self,
100
+ input_ids: Optional[torch.LongTensor] = None,
101
+ token_type_ids: Optional[torch.LongTensor] = None,
102
+ position_ids: Optional[torch.LongTensor] = None,
103
+ inputs_embeds: Optional[torch.FloatTensor] = None,
104
+ past_key_values_length: int = 0,
105
+ ) -> torch.Tensor:
106
+ if (input_ids is not None) == (inputs_embeds is not None):
107
+ raise ValueError('Must specify either input_ids or input_embeds!')
108
+ if input_ids is not None:
109
+ input_shape = input_ids.size()
110
+ else:
111
+ assert inputs_embeds is not None # just for type checking
112
+ input_shape = inputs_embeds.size()[:-1]
113
+
114
+ seq_length = input_shape[1]
115
+
116
+ if position_ids is None:
117
+ # great! ALiBi
118
+ pass
119
+
120
+ # Setting the token_type_ids to the registered buffer in constructor
121
+ # where it is all zeros, which usually occurs when it's auto-generated;
122
+ # registered buffer helps users when tracing the model without passing
123
+ # token_type_ids, solves issue #5664
124
+ if token_type_ids is None:
125
+ if hasattr(self, 'token_type_ids'):
126
+ assert isinstance(self.token_type_ids, torch.LongTensor)
127
+ buffered_token_type_ids = self.token_type_ids[:, :seq_length]
128
+ buffered_token_type_ids_expanded = buffered_token_type_ids.expand(
129
+ input_shape[0], seq_length)
130
+ token_type_ids = buffered_token_type_ids_expanded # type: ignore
131
+ else:
132
+ token_type_ids = torch.zeros(input_shape, # type: ignore
133
+ dtype=torch.long,
134
+ device=self.word_embeddings.device) # type: ignore # yapf: disable
135
+
136
+ if inputs_embeds is None:
137
+ inputs_embeds = self.word_embeddings(input_ids)
138
+ token_type_embeddings = self.token_type_embeddings(token_type_ids)
139
+
140
+ embeddings = inputs_embeds + token_type_embeddings
141
+ # no position embeddings! ALiBi
142
+ embeddings = self.LayerNorm(embeddings)
143
+ embeddings = self.dropout(embeddings)
144
+ return embeddings
145
+
146
+
147
+ class BertUnpadSelfAttention(nn.Module):
148
+ """Performs multi-headed self attention on a batch of unpadded sequences.
149
+
150
+ If Triton is installed, this module uses Flash Attention to greatly improve throughput.
151
+ The Flash Attention implementation used in Mosaic BERT supports arbitrary attention biases (which
152
+ we use to implement ALiBi), but does not support attention dropout. If either Triton is not installed
153
+ or `config.attention_probs_dropout_prob > 0`, the implementation will default to a
154
+ math-equivalent pytorch version, which is much slower.
155
+
156
+ See `forward` method for additional detail.
157
+ """
158
+
159
+ def __init__(self, config):
160
+ super().__init__()
161
+ if config.hidden_size % config.num_attention_heads != 0 and not hasattr(
162
+ config, 'embedding_size'):
163
+ raise ValueError(
164
+ f'The hidden size ({config.hidden_size}) is not a multiple of the number of attention '
165
+ f'heads ({config.num_attention_heads})')
166
+
167
+ self.num_attention_heads = config.num_attention_heads
168
+ self.attention_head_size = int(config.hidden_size /
169
+ config.num_attention_heads)
170
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
171
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
172
+ self.p_dropout = config.attention_probs_dropout_prob
173
+ self.Wqkv = nn.Linear(self.all_head_size, 3 * config.hidden_size)
174
+
175
+ # Warn if defaulting to pytorch because of import issues
176
+ if flash_attn_qkvpacked_func is None:
177
+ warnings.warn(
178
+ 'Unable to import Triton; defaulting MosaicBERT attention implementation to pytorch (this will reduce throughput when using this model).'
179
+ )
180
+
181
+ def forward(self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor,
182
+ max_seqlen_in_batch: int, indices: torch.Tensor,
183
+ attn_mask: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
184
+ """Perform self-attention.
185
+
186
+ If dropout is zero, then we can use the Triton kernel, so we do that. However, if not, we send through a standard PyTorch
187
+ implementation of self-attention.
188
+
189
+ The arguments are unpadded, and our implementations of attention require padded arguments,
190
+ so we first call `pad_input`. Once we compute attention, we re-unpad our outputs for the other layers.
191
+ The pad/unpad operations add overhead, but not sending pad tokens through ffs saves compute.
192
+ It is possible to write an unpadded implementation of attention (in Triton and PyTorch), which we will eventually do.
193
+
194
+ Args:
195
+ hidden_states: (total_nnz, dim)
196
+ cu_seqlens: (batch + 1,)
197
+ max_seqlen_in_batch: int
198
+ indices: (total_nnz,)
199
+ attn_mask: (batch, max_seqlen_in_batch)
200
+ bias: (batch, heads, max_seqlen_in_batch, max_seqlen_in_batch)
201
+
202
+ Returns:
203
+ attention: (total_nnz, dim)
204
+ """
205
+ qkv = self.Wqkv(hidden_states)
206
+ qkv = pad_input(qkv, indices, cu_seqlens.shape[0] - 1,
207
+ max_seqlen_in_batch) # batch, max_seqlen_in_batch, thd
208
+ qkv = rearrange(qkv,
209
+ 'b s (t h d) -> b s t h d',
210
+ t=3,
211
+ h=self.num_attention_heads)
212
+ if self.p_dropout or flash_attn_qkvpacked_func is None:
213
+ # if we have nonzero attention dropout (e.g. during fine-tuning) or no Triton, compute attention in PyTorch
214
+ q = qkv[:, :, 0, :, :].permute(0, 2, 1, 3) # b h s d
215
+ k = qkv[:, :, 1, :, :].permute(0, 2, 3, 1) # b h d s
216
+ v = qkv[:, :, 2, :, :].permute(0, 2, 1, 3) # b h s d
217
+ attention_scores = torch.matmul(q, k) / math.sqrt(
218
+ self.attention_head_size)
219
+ attention_scores = attention_scores + bias
220
+ attention_probs = nn.functional.softmax(attention_scores, dim=-1)
221
+ attention_probs = self.dropout(attention_probs)
222
+ attention = torch.matmul(attention_probs, v).permute(0, 2, 1,
223
+ 3) # b s h d
224
+ else:
225
+ # Triton implementation only supports 0 attention dropout
226
+ convert_dtype = qkv.dtype not in [torch.float16, torch.bfloat16]
227
+ if convert_dtype:
228
+ # Triton implementation only supports fp16 and bf16
229
+ orig_dtype = qkv.dtype
230
+ qkv = qkv.to(torch.float16)
231
+ bias_dtype = bias.dtype
232
+ bias = bias.to(torch.float16)
233
+ attention = flash_attn_qkvpacked_func(qkv, bias)
234
+ attention = attention.to(orig_dtype)
235
+ bias = bias.to(bias_dtype)
236
+ else:
237
+ attention = flash_attn_qkvpacked_func(qkv, bias)
238
+
239
+ # attn_mask is 1 for attend and 0 for don't
240
+ attention = unpad_input_only(attention, torch.squeeze(attn_mask) == 1)
241
+ return rearrange(attention, 'nnz h d -> nnz (h d)')
242
+
243
+
244
+ # Copy of transformer's library BertSelfOutput that will not be caught by surgery methods looking for HF BERT modules.
245
+ class BertSelfOutput(nn.Module):
246
+ """Computes the output of the attention layer.
247
+
248
+ This module is modeled after the Hugging Face BERT's
249
+ :class:`~transformers.model.bert.modeling_bert.BertSelfOutput`.
250
+ The implementation is identical. Rather than use the original module
251
+ directly, we re-implement it here so that Mosaic BERT's modules will not
252
+ be affected by any Composer surgery algorithm that modifies Hugging Face
253
+ BERT modules.
254
+ """
255
+
256
+ def __init__(self, config):
257
+ super().__init__()
258
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
259
+ self.LayerNorm = nn.LayerNorm(config.hidden_size,
260
+ eps=config.layer_norm_eps)
261
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
262
+
263
+ def forward(self, hidden_states: torch.Tensor,
264
+ input_tensor: torch.Tensor) -> torch.Tensor:
265
+ hidden_states = self.dense(hidden_states)
266
+ hidden_states = self.dropout(hidden_states)
267
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
268
+ return hidden_states
269
+
270
+
271
+ class BertUnpadAttention(nn.Module):
272
+ """Chains attention, Dropout, and LayerNorm for Mosaic BERT."""
273
+
274
+ def __init__(self, config):
275
+ super().__init__()
276
+ self.self = BertUnpadSelfAttention(config)
277
+ self.output = BertSelfOutput(config)
278
+
279
+ def forward(
280
+ self,
281
+ input_tensor: torch.Tensor,
282
+ cu_seqlens: torch.Tensor,
283
+ max_s: int,
284
+ subset_idx: Optional[torch.Tensor] = None,
285
+ indices: Optional[torch.Tensor] = None,
286
+ attn_mask: Optional[torch.Tensor] = None,
287
+ bias: Optional[torch.Tensor] = None,
288
+ ) -> torch.Tensor:
289
+ """Forward pass for scaled self-attention without padding.
290
+
291
+ Arguments:
292
+ input_tensor: (total_nnz, dim)
293
+ cu_seqlens: (batch + 1,)
294
+ max_s: int
295
+ subset_idx: () set of indices whose values we care about at the end of the layer
296
+ (e.g., the masked tokens, if this is the final layer).
297
+ indices: None or (total_nnz,)
298
+ attn_mask: None or (batch, max_seqlen_in_batch)
299
+ bias: None or (batch, heads, max_seqlen_in_batch, max_seqlen_in_batch)
300
+ """
301
+ self_output = self.self(input_tensor, cu_seqlens, max_s, indices,
302
+ attn_mask, bias)
303
+ if subset_idx is not None:
304
+ return self.output(index_first_axis(self_output, subset_idx),
305
+ index_first_axis(input_tensor, subset_idx))
306
+ else:
307
+ return self.output(self_output, input_tensor)
308
+
309
+
310
+ class BertGatedLinearUnitMLP(nn.Module):
311
+ """Applies the FFN at the end of each Mosaic BERT layer.
312
+
313
+ Compared to the default BERT architecture, this block replaces :class:`~transformers.model.bert.modeling_bert.BertIntermediate`
314
+ and :class:`~transformers.model.bert.modeling_bert.SelfOutput` with a single module that has similar functionality, but
315
+ introduces Gated Linear Units.
316
+
317
+ Note: Mosaic BERT adds parameters in order to implement Gated Linear Units. To keep parameter count consistent with that of a
318
+ standard Hugging Face BERT, scale down `config.intermediate_size` by 2/3. For example, a Mosaic BERT constructed with
319
+ `config.intermediate_size=2048` will have the same parameter footprint as its Hugging Face BERT counterpart constructed
320
+ with the `config.intermediate_size=3072`.
321
+ However, in most cases it will not be necessary to adjust `config.intermediate_size` since, despite the increased
322
+ parameter size, Mosaic BERT typically offers a net higher throughput than a Hugging Face BERT built from the same `config`.
323
+ """
324
+
325
+ def __init__(self, config):
326
+ super().__init__()
327
+ self.config = config
328
+ self.gated_layers = nn.Linear(config.hidden_size,
329
+ config.intermediate_size * 2,
330
+ bias=False)
331
+ self.act = nn.GELU(approximate='none')
332
+ self.wo = nn.Linear(config.intermediate_size, config.hidden_size)
333
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
334
+ self.layernorm = nn.LayerNorm(config.hidden_size,
335
+ eps=config.layer_norm_eps)
336
+
337
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
338
+ """Compute new hidden states from current hidden states.
339
+
340
+ Args:
341
+ hidden_states (torch.Tensor): The (unpadded) hidden states from
342
+ the attention layer [nnz, dim].
343
+ """
344
+ residual_connection = hidden_states
345
+ # compute the activation
346
+ hidden_states = self.gated_layers(hidden_states)
347
+ gated = hidden_states[:, :self.config.intermediate_size]
348
+ non_gated = hidden_states[:, self.config.intermediate_size:]
349
+ hidden_states = self.act(gated) * non_gated
350
+ hidden_states = self.dropout(hidden_states)
351
+ # multiply by the second matrix
352
+ hidden_states = self.wo(hidden_states)
353
+ # add the residual connection and post-LN
354
+ hidden_states = self.layernorm(hidden_states + residual_connection)
355
+ return hidden_states
356
+
357
+
358
+ class BertLayer(nn.Module):
359
+ """Composes the Mosaic BERT attention and FFN blocks into a single layer."""
360
+
361
+ def __init__(self, config):
362
+ super(BertLayer, self).__init__()
363
+ self.attention = BertUnpadAttention(config)
364
+ self.mlp = BertGatedLinearUnitMLP(config)
365
+
366
+ def forward(
367
+ self,
368
+ hidden_states: torch.Tensor,
369
+ cu_seqlens: torch.Tensor,
370
+ seqlen: int,
371
+ subset_idx: Optional[torch.Tensor] = None,
372
+ indices: Optional[torch.Tensor] = None,
373
+ attn_mask: Optional[torch.Tensor] = None,
374
+ bias: Optional[torch.Tensor] = None,
375
+ ) -> torch.Tensor:
376
+ """Forward pass for a BERT layer, including both attention and MLP.
377
+
378
+ Args:
379
+ hidden_states: (total_nnz, dim)
380
+ cu_seqlens: (batch + 1,)
381
+ seqlen: int
382
+ subset_idx: () set of indices whose values we care about at the end of the layer
383
+ (e.g., the masked tokens, if this is the final layer).
384
+ indices: None or (total_nnz,)
385
+ attn_mask: None or (batch, max_seqlen_in_batch)
386
+ bias: None or (batch, heads, max_seqlen_in_batch, max_seqlen_in_batch)
387
+ """
388
+ attention_output = self.attention(hidden_states, cu_seqlens, seqlen,
389
+ subset_idx, indices, attn_mask, bias)
390
+ layer_output = self.mlp(attention_output)
391
+ return layer_output
392
+
393
+
394
+ class BertEncoder(nn.Module):
395
+ """A stack of BERT layers providing the backbone of Mosaic BERT.
396
+
397
+ This module is modeled after the Hugging Face BERT's :class:`~transformers.model.bert.modeling_bert.BertEncoder`,
398
+ but with substantial modifications to implement unpadding and ALiBi.
399
+
400
+ Compared to the analogous Hugging Face BERT module, this module handles unpadding to reduce unnecessary computation
401
+ at padded tokens, and pre-computes attention biases to implement ALiBi.
402
+ """
403
+
404
+ def __init__(self, config):
405
+ super().__init__()
406
+ layer = BertLayer(config)
407
+ self.layer = nn.ModuleList(
408
+ [copy.deepcopy(layer) for _ in range(config.num_hidden_layers)])
409
+
410
+ self.num_attention_heads = config.num_attention_heads
411
+
412
+ # The alibi mask will be dynamically expanded if it is too small for
413
+ # the input the model receives. But it generally helps to initialize it
414
+ # to a reasonably large size to help pre-allocate CUDA memory.
415
+ # The default `alibi_starting_size` is 512.
416
+ self._current_alibi_size = int(config.alibi_starting_size)
417
+ self.alibi = torch.zeros(
418
+ (1, self.num_attention_heads, self._current_alibi_size,
419
+ self._current_alibi_size))
420
+ self.rebuild_alibi_tensor(size=config.alibi_starting_size)
421
+
422
+ def rebuild_alibi_tensor(self,
423
+ size: int,
424
+ device: Optional[Union[torch.device, str]] = None):
425
+ # Alibi
426
+ # Following https://github.com/ofirpress/attention_with_linear_biases/issues/5 (Implementation 1)
427
+ # In the causal case, you can exploit the fact that softmax is invariant to a uniform translation
428
+ # of the logits, which makes the math work out *after* applying causal masking. If no causal masking
429
+ # will be applied, it is necessary to construct the diagonal mask.
430
+ n_heads = self.num_attention_heads
431
+
432
+ def _get_alibi_head_slopes(n_heads: int) -> List[float]:
433
+
434
+ def get_slopes_power_of_2(n_heads: int) -> List[float]:
435
+ start = (2**(-2**-(math.log2(n_heads) - 3)))
436
+ ratio = start
437
+ return [start * ratio**i for i in range(n_heads)]
438
+
439
+ # In the paper, they only train models that have 2^a heads for some a. This function
440
+ # has some good properties that only occur when the input is a power of 2. To
441
+ # maintain that even when the number of heads is not a power of 2, we use a
442
+ # workaround.
443
+ if math.log2(n_heads).is_integer():
444
+ return get_slopes_power_of_2(n_heads)
445
+
446
+ closest_power_of_2 = 2**math.floor(math.log2(n_heads))
447
+ slopes_a = get_slopes_power_of_2(closest_power_of_2)
448
+ slopes_b = _get_alibi_head_slopes(2 * closest_power_of_2)
449
+ slopes_b = slopes_b[0::2][:n_heads - closest_power_of_2]
450
+ return slopes_a + slopes_b
451
+
452
+ context_position = torch.arange(size, device=device)[:, None]
453
+ memory_position = torch.arange(size, device=device)[None, :]
454
+ relative_position = torch.abs(memory_position - context_position)
455
+ # [n_heads, max_token_length, max_token_length]
456
+ relative_position = relative_position.unsqueeze(0).expand(
457
+ n_heads, -1, -1)
458
+ slopes = torch.Tensor(_get_alibi_head_slopes(n_heads)).to(device)
459
+ alibi = slopes.unsqueeze(1).unsqueeze(1) * -relative_position
460
+ # [1, n_heads, max_token_length, max_token_length]
461
+ alibi = alibi.unsqueeze(0)
462
+ assert alibi.shape == torch.Size([1, n_heads, size, size])
463
+
464
+ self._current_alibi_size = size
465
+ self.alibi = alibi
466
+
467
+ def forward(
468
+ self,
469
+ hidden_states: torch.Tensor,
470
+ attention_mask: torch.Tensor,
471
+ output_all_encoded_layers: Optional[bool] = True,
472
+ subset_mask: Optional[torch.Tensor] = None,
473
+ ) -> List[torch.Tensor]:
474
+
475
+ extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
476
+ extended_attention_mask = extended_attention_mask.to(
477
+ dtype=next(self.parameters()).dtype) # fp16 compatibility
478
+ extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
479
+
480
+ attention_mask_bool = attention_mask.bool()
481
+ batch, seqlen = hidden_states.shape[:2]
482
+ # Unpad inputs and mask. It will remove tokens that are padded.
483
+ # Assume ntokens is total number of tokens (padded and non-padded)
484
+ # and ntokens_unpad is total number of non-padded tokens.
485
+ # Then unpadding performs the following compression of the inputs:
486
+ # hidden_states[ntokens,hidden] -> hidden_states[ntokens_unpad,hidden]
487
+ hidden_states, indices, cu_seqlens, _ = unpad_input(
488
+ hidden_states, attention_mask_bool)
489
+
490
+ # Add alibi matrix to extended_attention_mask
491
+ if self._current_alibi_size < seqlen:
492
+ # Rebuild the alibi tensor when needed
493
+ warnings.warn(
494
+ f'Increasing alibi size from {self._current_alibi_size} to {seqlen}'
495
+ )
496
+ self.rebuild_alibi_tensor(size=seqlen, device=hidden_states.device)
497
+ elif self.alibi.device != hidden_states.device:
498
+ # Device catch-up
499
+ self.alibi = self.alibi.to(hidden_states.device)
500
+ alibi_bias = self.alibi[:, :, :seqlen, :seqlen]
501
+ attn_bias = extended_attention_mask[:, :, :seqlen, :seqlen]
502
+ alibi_attn_mask = attn_bias + alibi_bias
503
+
504
+ all_encoder_layers = []
505
+ if subset_mask is None:
506
+ for layer_module in self.layer:
507
+ hidden_states = layer_module(hidden_states,
508
+ cu_seqlens,
509
+ seqlen,
510
+ None,
511
+ indices,
512
+ attn_mask=attention_mask,
513
+ bias=alibi_attn_mask)
514
+ if output_all_encoded_layers:
515
+ all_encoder_layers.append(hidden_states)
516
+ # Pad inputs and mask. It will insert back zero-padded tokens.
517
+ # Assume ntokens is total number of tokens (padded and non-padded)
518
+ # and ntokens_unpad is total number of non-padded tokens.
519
+ # Then padding performs the following de-compression:
520
+ # hidden_states[ntokens_unpad,hidden] -> hidden_states[ntokens,hidden]
521
+ hidden_states = pad_input(hidden_states, indices, batch, seqlen)
522
+ else:
523
+ for i in range(len(self.layer) - 1):
524
+ layer_module = self.layer[i]
525
+ hidden_states = layer_module(hidden_states,
526
+ cu_seqlens,
527
+ seqlen,
528
+ None,
529
+ indices,
530
+ attn_mask=attention_mask,
531
+ bias=alibi_attn_mask)
532
+ if output_all_encoded_layers:
533
+ all_encoder_layers.append(hidden_states)
534
+ subset_idx = torch.nonzero(subset_mask[attention_mask_bool],
535
+ as_tuple=False).flatten()
536
+ hidden_states = self.layer[-1](hidden_states,
537
+ cu_seqlens,
538
+ seqlen,
539
+ subset_idx=subset_idx,
540
+ indices=indices,
541
+ attn_mask=attention_mask,
542
+ bias=alibi_attn_mask)
543
+
544
+ if not output_all_encoded_layers:
545
+ all_encoder_layers.append(hidden_states)
546
+ return all_encoder_layers
547
+
548
+
549
+ class BertPooler(nn.Module):
550
+
551
+ def __init__(self, config):
552
+ super(BertPooler, self).__init__()
553
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
554
+ self.activation = nn.Tanh()
555
+
556
+ def forward(self,
557
+ hidden_states: torch.Tensor,
558
+ pool: Optional[bool] = True) -> torch.Tensor:
559
+ # We "pool" the model by simply taking the hidden state corresponding
560
+ # to the first token.
561
+ first_token_tensor = hidden_states[:, 0] if pool else hidden_states
562
+ pooled_output = self.dense(first_token_tensor)
563
+ pooled_output = self.activation(pooled_output)
564
+ return pooled_output
565
+
566
+
567
+ class BertPredictionHeadTransform(nn.Module):
568
+
569
+ def __init__(self, config):
570
+ super().__init__()
571
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
572
+ if isinstance(config.hidden_act, str):
573
+ self.transform_act_fn = ACT2FN[config.hidden_act]
574
+ else:
575
+ self.transform_act_fn = config.hidden_act
576
+ self.LayerNorm = torch.nn.LayerNorm(config.hidden_size, eps=1e-12)
577
+
578
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
579
+ hidden_states = self.dense(hidden_states)
580
+ hidden_states = self.transform_act_fn(hidden_states)
581
+ hidden_states = self.LayerNorm(hidden_states)
582
+ return hidden_states
583
+
584
+
585
+ class BertModel(BertPreTrainedModel):
586
+ """Overall BERT model.
587
+
588
+ Args:
589
+ config: a BertConfig class instance with the configuration to build a new model
590
+
591
+ Inputs:
592
+ `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length]
593
+ with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts
594
+ `extract_features.py`, `run_classifier.py` and `run_squad.py`)
595
+ `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token
596
+ types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to
597
+ a `sentence B` token (see BERT paper for more details).
598
+ `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices
599
+ selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
600
+ input sequence length in the current batch. It's the mask that we typically use for attention when
601
+ a batch has varying length sentences.
602
+ `output_all_encoded_layers`: boolean which controls the content of the `encoded_layers` output as described below. Default: `True`.
603
+
604
+ Outputs: Tuple of (encoded_layers, pooled_output)
605
+ `encoded_layers`: controlled by `output_all_encoded_layers` argument:
606
+ - `output_all_encoded_layers=True`: outputs a list of the full sequences of encoded-hidden-states at the end
607
+ of each attention block (i.e. 12 full sequences for BERT-base, 24 for BERT-large), each
608
+ encoded-hidden-state is a torch.FloatTensor of size [batch_size, sequence_length, hidden_size],
609
+ - `output_all_encoded_layers=False`: outputs only the full sequence of hidden-states corresponding
610
+ to the last attention block of shape [batch_size, sequence_length, hidden_size],
611
+ `pooled_output`: a torch.FloatTensor of size [batch_size, hidden_size] which is the output of a
612
+ classifier pretrained on top of the hidden state associated to the first character of the
613
+ input (`CLS`) to train on the Next-Sentence task (see BERT's paper).
614
+
615
+ Example usage:
616
+ ```python
617
+ # Already been converted into WordPiece token ids
618
+ input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
619
+ input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
620
+ token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])
621
+ config = modeling.BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
622
+ num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)
623
+ model = BertModel(config=config)
624
+ all_encoder_layers, pooled_output = model(input_ids, token_type_ids, input_mask)
625
+ ```
626
+ """
627
+
628
+ def __init__(self, config, add_pooling_layer=True):
629
+ super(BertModel, self).__init__(config)
630
+ self.embeddings = BertEmbeddings(config)
631
+ self.encoder = BertEncoder(config)
632
+ self.pooler = BertPooler(config) if add_pooling_layer else None
633
+ self.post_init()
634
+
635
+ def get_input_embeddings(self):
636
+ return self.embeddings.word_embeddings
637
+
638
+ def set_input_embeddings(self, value):
639
+ self.embeddings.word_embeddings = value
640
+
641
+ def forward(
642
+ self,
643
+ input_ids: torch.Tensor,
644
+ token_type_ids: Optional[torch.Tensor] = None,
645
+ attention_mask: Optional[torch.Tensor] = None,
646
+ position_ids: Optional[torch.Tensor] = None,
647
+ output_all_encoded_layers: Optional[bool] = False,
648
+ masked_tokens_mask: Optional[torch.Tensor] = None,
649
+ **kwargs
650
+ ) -> Tuple[Union[List[torch.Tensor], torch.Tensor], Optional[torch.Tensor]]:
651
+ if attention_mask is None:
652
+ attention_mask = torch.ones_like(input_ids)
653
+ if token_type_ids is None:
654
+ token_type_ids = torch.zeros_like(input_ids)
655
+
656
+ embedding_output = self.embeddings(input_ids, token_type_ids,
657
+ position_ids)
658
+
659
+ subset_mask = []
660
+ first_col_mask = []
661
+
662
+ if masked_tokens_mask is None:
663
+ subset_mask = None
664
+ else:
665
+ first_col_mask = torch.zeros_like(masked_tokens_mask)
666
+ first_col_mask[:, 0] = True
667
+ subset_mask = masked_tokens_mask | first_col_mask
668
+
669
+ encoder_outputs = self.encoder(
670
+ embedding_output,
671
+ attention_mask,
672
+ output_all_encoded_layers=output_all_encoded_layers,
673
+ subset_mask=subset_mask)
674
+
675
+ if masked_tokens_mask is None:
676
+ sequence_output = encoder_outputs[-1]
677
+ pooled_output = self.pooler(
678
+ sequence_output) if self.pooler is not None else None
679
+ else:
680
+ # TD [2022-03-01]: the indexing here is very tricky.
681
+ attention_mask_bool = attention_mask.bool()
682
+ subset_idx = subset_mask[attention_mask_bool] # type: ignore
683
+ sequence_output = encoder_outputs[-1][
684
+ masked_tokens_mask[attention_mask_bool][subset_idx]]
685
+ if self.pooler is not None:
686
+ pool_input = encoder_outputs[-1][
687
+ first_col_mask[attention_mask_bool][subset_idx]]
688
+ pooled_output = self.pooler(pool_input, pool=False)
689
+ else:
690
+ pooled_output = None
691
+
692
+ if not output_all_encoded_layers:
693
+ encoder_outputs = sequence_output
694
+
695
+ if self.pooler is not None:
696
+ return encoder_outputs, pooled_output
697
+
698
+ return encoder_outputs, None
699
+
700
+
701
+ ###################
702
+ # Bert Heads
703
+ ###################
704
+ class BertLMPredictionHead(nn.Module):
705
+
706
+ def __init__(self, config, bert_model_embedding_weights):
707
+ super().__init__()
708
+ self.transform = BertPredictionHeadTransform(config)
709
+ # The output weights are the same as the input embeddings, but there is
710
+ # an output-only bias for each token.
711
+ self.decoder = nn.Linear(bert_model_embedding_weights.size(1),
712
+ bert_model_embedding_weights.size(0))
713
+ self.decoder.weight = bert_model_embedding_weights
714
+
715
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
716
+ hidden_states = self.transform(hidden_states)
717
+ hidden_states = self.decoder(hidden_states)
718
+ return hidden_states
719
+
720
+
721
+ class BertOnlyMLMHead(nn.Module):
722
+
723
+ def __init__(self, config, bert_model_embedding_weights):
724
+ super().__init__()
725
+ self.predictions = BertLMPredictionHead(config,
726
+ bert_model_embedding_weights)
727
+
728
+ def forward(self, sequence_output: torch.Tensor) -> torch.Tensor:
729
+ prediction_scores = self.predictions(sequence_output)
730
+ return prediction_scores
731
+
732
+
733
+ class BertOnlyNSPHead(nn.Module):
734
+
735
+ def __init__(self, config):
736
+ super().__init__()
737
+ self.seq_relationship = nn.Linear(config.hidden_size, 2)
738
+
739
+ def forward(self, pooled_output: torch.Tensor) -> torch.Tensor:
740
+ seq_relationship_score = self.seq_relationship(pooled_output)
741
+ return seq_relationship_score
742
+
743
+
744
+ #####################
745
+ # Various Bert models
746
+ #####################
747
+
748
+
749
+ class BertForPreTraining(BertPreTrainedModel):
750
+ #TBD: Coming in Future Commit
751
+ pass
752
+
753
+
754
+ class BertLMHeadModel(BertPreTrainedModel):
755
+ #TBD: Coming in Future Commit
756
+ pass
757
+
758
+
759
+ class BertForMaskedLM(BertPreTrainedModel):
760
+
761
+ def __init__(self, config):
762
+ super().__init__(config)
763
+
764
+ if config.is_decoder:
765
+ warnings.warn(
766
+ 'If you want to use `BertForMaskedLM` make sure `config.is_decoder=False` for '
767
+ 'bi-directional self-attention.')
768
+
769
+ self.bert = BertModel(config, add_pooling_layer=False)
770
+ self.cls = BertOnlyMLMHead(config,
771
+ self.bert.embeddings.word_embeddings.weight)
772
+
773
+ # Initialize weights and apply final processing
774
+ self.post_init()
775
+
776
+ @classmethod
777
+ def from_composer(cls,
778
+ pretrained_checkpoint,
779
+ state_dict=None,
780
+ cache_dir=None,
781
+ from_tf=False,
782
+ config=None,
783
+ *inputs,
784
+ **kwargs):
785
+ """Load from pre-trained."""
786
+ model = cls(config, *inputs, **kwargs)
787
+ if from_tf:
788
+ raise ValueError(
789
+ 'Mosaic BERT does not support loading TensorFlow weights.')
790
+
791
+ state_dict = torch.load(pretrained_checkpoint)
792
+ # If the state_dict was saved after wrapping with `composer.HuggingFaceModel`, it takes on the `model` prefix
793
+ consume_prefix_in_state_dict_if_present(state_dict, prefix='model.')
794
+ missing_keys, unexpected_keys = model.load_state_dict(state_dict,
795
+ strict=False)
796
+
797
+ if len(missing_keys) > 0:
798
+ logger.warning(
799
+ f"Found these missing keys in the checkpoint: {', '.join(missing_keys)}"
800
+ )
801
+ if len(unexpected_keys) > 0:
802
+ logger.warning(
803
+ f"Found these unexpected keys in the checkpoint: {', '.join(unexpected_keys)}"
804
+ )
805
+
806
+ return model
807
+
808
+ def get_output_embeddings(self):
809
+ return self.cls.predictions.decoder
810
+
811
+ def set_output_embeddings(self, new_embeddings):
812
+ self.cls.predictions.decoder = new_embeddings
813
+
814
+ def forward(
815
+ self,
816
+ input_ids: Optional[torch.Tensor] = None,
817
+ attention_mask: Optional[torch.Tensor] = None,
818
+ token_type_ids: Optional[torch.Tensor] = None,
819
+ position_ids: Optional[torch.Tensor] = None,
820
+ head_mask: Optional[torch.Tensor] = None,
821
+ inputs_embeds: Optional[torch.Tensor] = None,
822
+ encoder_hidden_states: Optional[torch.Tensor] = None,
823
+ encoder_attention_mask: Optional[torch.Tensor] = None,
824
+ labels: Optional[torch.Tensor] = None,
825
+ output_attentions: Optional[bool] = None,
826
+ output_hidden_states: Optional[bool] = None,
827
+ return_dict: Optional[bool] = None,
828
+ ) -> Union[Tuple[torch.Tensor], MaskedLMOutput]:
829
+ # labels should be a `torch.LongTensor` of shape
830
+ # `(batch_size, sequence_length)`. These are used for computing the
831
+ # masked language modeling loss.
832
+ #
833
+ # Indices should be in `[-100, 0, ..., config.vocab_size]` (see
834
+ # `input_ids` docstring) Tokens with indices set to `-100` are ignored
835
+ # (masked), the loss is only computed for the tokens with labels in `[0,
836
+ # ..., config.vocab_size]`
837
+ #
838
+ # Prediction scores are only computed for masked tokens and the (bs,
839
+ # seqlen) dimensions are flattened
840
+ if (input_ids is not None) == (inputs_embeds is not None):
841
+ raise ValueError('Must specify either input_ids or input_embeds!')
842
+
843
+ if labels is None:
844
+ masked_tokens_mask = None
845
+ else:
846
+ masked_tokens_mask = labels > 0
847
+
848
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
849
+
850
+ outputs = self.bert(
851
+ input_ids,
852
+ attention_mask=attention_mask,
853
+ token_type_ids=token_type_ids,
854
+ position_ids=position_ids,
855
+ head_mask=head_mask,
856
+ inputs_embeds=inputs_embeds,
857
+ encoder_hidden_states=encoder_hidden_states,
858
+ encoder_attention_mask=encoder_attention_mask,
859
+ output_attentions=output_attentions,
860
+ output_hidden_states=output_hidden_states,
861
+ return_dict=return_dict,
862
+ masked_tokens_mask=masked_tokens_mask,
863
+ )
864
+
865
+ sequence_output = outputs[0]
866
+ prediction_scores = self.cls(sequence_output)
867
+
868
+ loss = None
869
+ if labels is not None:
870
+ # Compute loss
871
+ loss_fct = nn.CrossEntropyLoss()
872
+ masked_token_idx = torch.nonzero(labels.flatten() > 0,
873
+ as_tuple=False).flatten()
874
+ loss = loss_fct(prediction_scores,
875
+ labels.flatten()[masked_token_idx])
876
+
877
+ assert input_ids is not None, 'Coding error; please open an issue'
878
+ batch, seqlen = input_ids.shape[:2]
879
+ prediction_scores = rearrange(index_put_first_axis(
880
+ prediction_scores, masked_token_idx, batch * seqlen),
881
+ '(b s) d -> b s d',
882
+ b=batch)
883
+
884
+ if not return_dict:
885
+ output = (prediction_scores,) + outputs[2:]
886
+ return ((loss,) + output) if loss is not None else output
887
+
888
+ return MaskedLMOutput(
889
+ loss=loss,
890
+ logits=prediction_scores,
891
+ hidden_states=None,
892
+ attentions=None,
893
+ )
894
+
895
+ def prepare_inputs_for_generation(self, input_ids: torch.Tensor,
896
+ attention_mask: torch.Tensor,
897
+ **model_kwargs):
898
+ input_shape = input_ids.shape
899
+ effective_batch_size = input_shape[0]
900
+
901
+ # add a dummy token
902
+ if self.config.pad_token_id is None:
903
+ raise ValueError('The PAD token should be defined for generation')
904
+
905
+ attention_mask = torch.cat([
906
+ attention_mask,
907
+ attention_mask.new_zeros((attention_mask.shape[0], 1))
908
+ ],
909
+ dim=-1)
910
+ dummy_token = torch.full((effective_batch_size, 1),
911
+ self.config.pad_token_id,
912
+ dtype=torch.long,
913
+ device=input_ids.device)
914
+ input_ids = torch.cat([input_ids, dummy_token], dim=1)
915
+
916
+ return {'input_ids': input_ids, 'attention_mask': attention_mask}
917
+
918
+
919
+ class BertForNextSentencePrediction(BertPreTrainedModel):
920
+ #TBD: Push in future commit
921
+ pass
922
+
923
+
924
+ class BertForSequenceClassification(BertPreTrainedModel):
925
+ """Bert Model transformer with a sequence classification/regression head.
926
+
927
+ This head is just a linear layer on top of the pooled output. Used for,
928
+ e.g., GLUE tasks.
929
+ """
930
+
931
+ def __init__(self, config):
932
+ super().__init__(config)
933
+ self.num_labels = config.num_labels
934
+ self.config = config
935
+
936
+ self.bert = BertModel(config)
937
+ classifier_dropout = (config.classifier_dropout
938
+ if config.classifier_dropout is not None else
939
+ config.hidden_dropout_prob)
940
+ self.dropout = nn.Dropout(classifier_dropout)
941
+ self.classifier = nn.Linear(config.hidden_size, config.num_labels)
942
+
943
+ # Initialize weights and apply final processing
944
+ self.post_init()
945
+
946
+ @classmethod
947
+ def from_composer(cls,
948
+ pretrained_checkpoint,
949
+ state_dict=None,
950
+ cache_dir=None,
951
+ from_tf=False,
952
+ config=None,
953
+ *inputs,
954
+ **kwargs):
955
+ """Load from pre-trained."""
956
+ model = cls(config, *inputs, **kwargs)
957
+ if from_tf:
958
+ raise ValueError(
959
+ 'Mosaic BERT does not support loading TensorFlow weights.')
960
+
961
+ state_dict = torch.load(pretrained_checkpoint)
962
+ # If the state_dict was saved after wrapping with `composer.HuggingFaceModel`, it takes on the `model` prefix
963
+ consume_prefix_in_state_dict_if_present(state_dict, prefix='model.')
964
+ missing_keys, unexpected_keys = model.load_state_dict(state_dict,
965
+ strict=False)
966
+
967
+ if len(missing_keys) > 0:
968
+ logger.warning(
969
+ f"Found these missing keys in the checkpoint: {', '.join(missing_keys)}"
970
+ )
971
+ if len(unexpected_keys) > 0:
972
+ logger.warning(
973
+ f"Found these unexpected keys in the checkpoint: {', '.join(unexpected_keys)}"
974
+ )
975
+
976
+ return model
977
+
978
+ def forward(
979
+ self,
980
+ input_ids: Optional[torch.Tensor] = None,
981
+ attention_mask: Optional[torch.Tensor] = None,
982
+ token_type_ids: Optional[torch.Tensor] = None,
983
+ position_ids: Optional[torch.Tensor] = None,
984
+ head_mask: Optional[torch.Tensor] = None,
985
+ inputs_embeds: Optional[torch.Tensor] = None,
986
+ labels: Optional[torch.Tensor] = None,
987
+ output_attentions: Optional[bool] = None,
988
+ output_hidden_states: Optional[bool] = None,
989
+ return_dict: Optional[bool] = None,
990
+ ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:
991
+ # labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
992
+ # Labels for computing the sequence classification/regression loss.
993
+ # Indices should be in `[0, ..., config.num_labels - 1]`.
994
+ # If `config.num_labels == 1` a regression loss is computed
995
+ # (mean-square loss). If `config.num_labels > 1` a classification loss
996
+ # is computed (cross-entropy).
997
+
998
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
999
+
1000
+ outputs = self.bert(
1001
+ input_ids,
1002
+ attention_mask=attention_mask,
1003
+ token_type_ids=token_type_ids,
1004
+ position_ids=position_ids,
1005
+ head_mask=head_mask,
1006
+ inputs_embeds=inputs_embeds,
1007
+ output_attentions=output_attentions,
1008
+ output_hidden_states=output_hidden_states,
1009
+ return_dict=return_dict,
1010
+ )
1011
+
1012
+ pooled_output = outputs[1]
1013
+
1014
+ pooled_output = self.dropout(pooled_output)
1015
+ logits = self.classifier(pooled_output)
1016
+
1017
+ loss = None
1018
+ if labels is not None:
1019
+ # Compute loss
1020
+ if self.config.problem_type is None:
1021
+ if self.num_labels == 1:
1022
+ self.config.problem_type = 'regression'
1023
+ elif self.num_labels > 1 and (labels.dtype == torch.long or
1024
+ labels.dtype == torch.int):
1025
+ self.config.problem_type = 'single_label_classification'
1026
+ else:
1027
+ self.config.problem_type = 'multi_label_classification'
1028
+
1029
+ if self.config.problem_type == 'regression':
1030
+ loss_fct = nn.MSELoss()
1031
+ if self.num_labels == 1:
1032
+ loss = loss_fct(logits.squeeze(), labels.squeeze())
1033
+ else:
1034
+ loss = loss_fct(logits, labels)
1035
+ elif self.config.problem_type == 'single_label_classification':
1036
+ loss_fct = nn.CrossEntropyLoss()
1037
+ loss = loss_fct(logits.view(-1, self.num_labels),
1038
+ labels.view(-1))
1039
+ elif self.config.problem_type == 'multi_label_classification':
1040
+ loss_fct = nn.BCEWithLogitsLoss()
1041
+ loss = loss_fct(logits, labels)
1042
+
1043
+ if not return_dict:
1044
+ output = (logits,) + outputs[2:]
1045
+ return ((loss,) + output) if loss is not None else output
1046
+
1047
+ return SequenceClassifierOutput(
1048
+ loss=loss,
1049
+ logits=logits,
1050
+ hidden_states=None,
1051
+ attentions=None,
1052
+ )
1053
+
1054
+
1055
+ class BertForMultipleChoice(BertPreTrainedModel):
1056
+ #TBD: Push in future commit
1057
+ pass
1058
+
1059
+
1060
+ class BertForTokenClassification(BertPreTrainedModel):
1061
+ #TBD: Push in future commit
1062
+ pass
1063
+
1064
+
1065
+ class BertForQuestionAnswering(BertPreTrainedModel):
1066
+ """Bert Model with a span classification head.
1067
+
1068
+ This is used for extractive question-answering tasks like SQuAD (a linear
1069
+ layers on top of the hidden states' output to compute `span start logits`
1070
+ and `span end logits`).
1071
+ """
1072
+ #TBD: Push in future commit
bert_padding.py ADDED
@@ -0,0 +1,159 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ """Helper functions for padding and unpadding batches.
8
+
9
+ These functions are used extensively throughout the Mosaic BERT implementation
10
+ in `bert_layers.py`.
11
+ """
12
+
13
+ from typing import Tuple, cast
14
+
15
+ import torch
16
+ import torch.nn.functional as F
17
+ from einops import rearrange, repeat
18
+
19
+
20
+ class IndexFirstAxis(torch.autograd.Function):
21
+
22
+ @staticmethod
23
+ def forward(ctx, input: torch.Tensor,
24
+ indices: torch.Tensor) -> torch.Tensor:
25
+ """Get just the values of `input` which are at `indices`.
26
+
27
+ Arguments:
28
+ ctx: the autograd context object
29
+ input: (b, ...) 2+ dimensional tensor
30
+ indices: (num_idx) 1D tensor
31
+ """
32
+ ctx.save_for_backward(indices)
33
+ assert input.ndim >= 2
34
+ ctx.first_axis_dim, other_shape = input.shape[0], input.shape[
35
+ 1:] # type: ignore
36
+ second_dim = other_shape.numel(
37
+ ) # product of sizes of all but first dimension
38
+ # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing.
39
+ return torch.gather(
40
+ rearrange(input, 'b ... -> b (...)'), # (b, ...) -> (b, second_dim)
41
+ 0,
42
+ repeat(indices, 'z -> z d',
43
+ d=second_dim) # (indices,) -> (indices, second_dim)
44
+ ).reshape(-1, *other_shape) # (num_idx, ...)
45
+
46
+ @staticmethod
47
+ def backward(ctx, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None]:
48
+ indices, = ctx.saved_tensors
49
+ assert grad_output.ndim >= 2
50
+ other_shape = grad_output.shape[1:]
51
+ grad_output = rearrange(grad_output, 'b ... -> b (...)')
52
+ grad_input = torch.zeros([ctx.first_axis_dim, grad_output.shape[1]],
53
+ device=grad_output.device,
54
+ dtype=grad_output.dtype)
55
+ # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing.
56
+ # grad_input[indices] = grad_output
57
+ grad_input.scatter_(0,
58
+ repeat(indices, 'z -> z d', d=grad_output.shape[1]),
59
+ grad_output)
60
+ return grad_input.reshape(ctx.first_axis_dim, *other_shape), None
61
+
62
+
63
+ index_first_axis = IndexFirstAxis.apply
64
+
65
+
66
+ class IndexPutFirstAxis(torch.autograd.Function):
67
+
68
+ @staticmethod
69
+ def forward(ctx, values: torch.Tensor, indices: torch.Tensor,
70
+ first_axis_dim) -> torch.Tensor:
71
+ ctx.save_for_backward(indices)
72
+ assert indices.ndim == 1
73
+ assert values.ndim >= 2
74
+ output = torch.zeros(first_axis_dim,
75
+ *values.shape[1:],
76
+ device=values.device,
77
+ dtype=values.dtype)
78
+ output[indices] = values
79
+ return output
80
+
81
+ @staticmethod
82
+ def backward(ctx,
83
+ grad_output: torch.Tensor) -> Tuple[torch.Tensor, None, None]:
84
+ indices, = ctx.saved_tensors
85
+ grad_values = grad_output[indices]
86
+ return grad_values, None, None
87
+
88
+
89
+ index_put_first_axis = IndexPutFirstAxis.apply
90
+
91
+
92
+ def unpad_input(
93
+ hidden_states: torch.Tensor,
94
+ attention_mask: torch.Tensor,
95
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
96
+ """Remove padding from input sequences.
97
+
98
+ Arguments:
99
+ hidden_states: (batch, seqlen, ...)
100
+ attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid.
101
+
102
+ Returns:
103
+ hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
104
+ indices: (total_nnz)
105
+ cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states.
106
+ max_seqlen_in_batch: int ()
107
+ """
108
+ seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
109
+ indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
110
+ max_seqlen_in_batch = int(seqlens_in_batch.max().item())
111
+ cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32),
112
+ (1, 0))
113
+ # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the
114
+ # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim
115
+ # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to
116
+ # index with integer indices. Moreover, torch's index is a bit slower than it needs to be,
117
+ # so we write custom forward and backward to make it a bit faster.
118
+ hidden_states = cast(
119
+ torch.Tensor,
120
+ index_first_axis(rearrange(hidden_states, 'b s ... -> (b s) ...'),
121
+ indices))
122
+ return hidden_states, indices, cu_seqlens, max_seqlen_in_batch
123
+
124
+
125
+ def unpad_input_only(
126
+ hidden_states: torch.Tensor,
127
+ attention_mask: torch.Tensor,
128
+ ) -> torch.Tensor:
129
+ """Like unpad_input, but only return the unpadded first tensor.
130
+
131
+ Save a small amount of overhead.
132
+
133
+ Arguments:
134
+ hidden_states: (batch, seqlen, ...)
135
+ attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid.
136
+
137
+ Returns:
138
+ hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
139
+ """
140
+ indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
141
+ return index_first_axis(rearrange(hidden_states, 'b s ... -> (b s) ...'),
142
+ indices)
143
+
144
+
145
+ def pad_input(hidden_states: torch.Tensor, indices: torch.Tensor, batch: int,
146
+ seqlen: int) -> torch.Tensor:
147
+ """Add padding to sequences.
148
+
149
+ Arguments:
150
+ hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
151
+ indices: (total_nnz)
152
+ batch: int batch_size
153
+ seqlen: int max sequence length
154
+
155
+ Returns:
156
+ hidden_states: (batch, seqlen, ...)
157
+ """
158
+ output = index_put_first_axis(hidden_states, indices, batch * seqlen)
159
+ return rearrange(output, '(b s) ... -> b s ...', b=batch)
config.json ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "mosaicml/mosaic-bert-base-seqlen-256",
3
+ "alibi_starting_size": 512,
4
+ "architectures": [
5
+ "BertForMaskedLM"
6
+ ],
7
+ "attention_probs_dropout_prob": 0.0,
8
+ "auto_map": {
9
+ "AutoConfig": "configuration_bert.BertConfig",
10
+ "AutoModelForMaskedLM": "bert_layers.BertForMaskedLM"
11
+ },
12
+ "classifier_dropout": null,
13
+ "gradient_checkpointing": false,
14
+ "hidden_act": "gelu",
15
+ "hidden_dropout_prob": 0.1,
16
+ "hidden_size": 768,
17
+ "initializer_range": 0.02,
18
+ "intermediate_size": 3072,
19
+ "layer_norm_eps": 1e-12,
20
+ "max_position_embeddings": 512,
21
+ "model_type": "bert",
22
+ "num_attention_heads": 12,
23
+ "num_hidden_layers": 12,
24
+ "pad_token_id": 0,
25
+ "position_embedding_type": "absolute",
26
+ "torch_dtype": "float32",
27
+ "transformers_version": "4.25.1",
28
+ "type_vocab_size": 2,
29
+ "use_cache": true,
30
+ "vocab_size": 30528
31
+ }
configuration_bert.py ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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__(attention_probs_dropout_prob=attention_probs_dropout_prob, **kwargs)
25
+ 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
pytorch_model.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:a81e58ef4b51cc36e43cea855a99f699d9b6056014489b37d79f0917e616920c
3
+ size 549647297