Feature Extraction
Transformers
Safetensors
English
bamboo
custom_code
Yixin Song commited on
Commit
da6a853
1 Parent(s): ed33910
Files changed (1) hide show
  1. modeling_attn_mask_utils.py +0 -500
modeling_attn_mask_utils.py DELETED
@@ -1,500 +0,0 @@
1
- # Copyright 2023 The HuggingFace Team. All rights reserved.
2
- #
3
- # Licensed under the Apache License, Version 2.0 (the "License");
4
- # you may not use this file except in compliance with the License.
5
- # You may obtain a copy of the License at
6
- #
7
- # http://www.apache.org/licenses/LICENSE-2.0
8
- #
9
- # Unless required by applicable law or agreed to in writing, software
10
- # distributed under the License is distributed on an "AS IS" BASIS,
11
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
- # See the License for the specific language governing permissions and
13
- # limitations under the License.
14
- from dataclasses import dataclass
15
- from typing import List, Optional, Tuple, Union
16
-
17
- import torch
18
-
19
-
20
- @dataclass
21
- class AttentionMaskConverter:
22
- """
23
- A utility attention mask class that allows one to:
24
- - Create a causal 4d mask
25
- - Create a causal 4d mask with slided window
26
- - Convert a 2d attention mask (batch_size, query_length) to a 4d attention mask (batch_size, 1, query_length,
27
- key_value_length) that can be multiplied with attention scores
28
-
29
- Examples:
30
-
31
- ```python
32
- >>> import torch
33
- >>> from transformers.modeling_attn_mask_utils import AttentionMaskConverter
34
-
35
- >>> converter = AttentionMaskConverter(True)
36
- >>> converter.to_4d(torch.tensor([[0, 0, 0, 1, 1]]), 5, key_value_length=5, dtype=torch.float32)
37
- tensor([[[[-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38],
38
- [-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38],
39
- [-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38],
40
- [-3.4028e+38, -3.4028e+38, -3.4028e+38, 0.0000e+00, -3.4028e+38],
41
- [-3.4028e+38, -3.4028e+38, -3.4028e+38, 0.0000e+00, 0.0000e+00]]]])
42
- ```
43
-
44
- Parameters:
45
- is_causal (`bool`):
46
- Whether the attention mask should be a uni-directional (causal) or bi-directional mask.
47
-
48
- sliding_window (`int`, *optional*):
49
- Optionally, the sliding window masks can be created if `sliding_window` is defined to a positive integer.
50
- """
51
-
52
- is_causal: bool
53
- sliding_window: int
54
-
55
- def __init__(self, is_causal: bool, sliding_window: Optional[int] = None):
56
- self.is_causal = is_causal
57
- self.sliding_window = sliding_window
58
-
59
- if self.sliding_window is not None and self.sliding_window <= 0:
60
- raise ValueError(
61
- f"Make sure that when passing `sliding_window` that its value is a strictly positive integer, not `{self.sliding_window}`"
62
- )
63
-
64
- def to_causal_4d(
65
- self,
66
- batch_size: int,
67
- query_length: int,
68
- key_value_length: int,
69
- dtype: torch.dtype,
70
- device: Union[torch.device, "str"] = "cpu",
71
- ) -> Optional[torch.Tensor]:
72
- """
73
- Creates a causal 4D mask of (bsz, head_dim=1, query_length, key_value_length) shape and adds large negative
74
- bias to upper right hand triangular matrix (causal mask).
75
- """
76
- if not self.is_causal:
77
- raise ValueError(f"Please use `to_causal_4d` only if {self.__class__} has `is_causal` set to True.")
78
-
79
- # If shape is not cached, create a new causal mask and cache it
80
- input_shape = (batch_size, query_length)
81
- past_key_values_length = key_value_length - query_length
82
-
83
- # create causal mask
84
- # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
85
- causal_4d_mask = None
86
- if input_shape[-1] > 1 or self.sliding_window is not None:
87
- causal_4d_mask = self._make_causal_mask(
88
- input_shape,
89
- dtype,
90
- device=device,
91
- past_key_values_length=past_key_values_length,
92
- sliding_window=self.sliding_window,
93
- )
94
-
95
- return causal_4d_mask
96
-
97
- def to_4d(
98
- self,
99
- attention_mask_2d: torch.Tensor,
100
- query_length: int,
101
- dtype: torch.dtype,
102
- key_value_length: Optional[int] = None,
103
- ) -> torch.Tensor:
104
- """
105
- Converts 2D attention mask to 4D attention mask by expanding mask to (bsz, head_dim=1, query_length,
106
- key_value_length) shape and by adding a large negative bias to not-attended positions. If attention_mask is
107
- causal, a causal mask will be added.
108
- """
109
- input_shape = (attention_mask_2d.shape[0], query_length)
110
-
111
- # create causal mask
112
- # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
113
- causal_4d_mask = None
114
- if (input_shape[-1] > 1 or self.sliding_window is not None) and self.is_causal:
115
- if key_value_length is None:
116
- raise ValueError(
117
- "This attention mask converter is causal. Make sure to pass `key_value_length` to correctly create a causal mask."
118
- )
119
-
120
- past_key_values_length = key_value_length - query_length
121
- causal_4d_mask = self._make_causal_mask(
122
- input_shape,
123
- dtype,
124
- device=attention_mask_2d.device,
125
- past_key_values_length=past_key_values_length,
126
- sliding_window=self.sliding_window,
127
- )
128
- elif self.sliding_window is not None:
129
- raise NotImplementedError("Sliding window is currently only implemented for causal masking")
130
-
131
- # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
132
- expanded_attn_mask = self._expand_mask(attention_mask_2d, dtype, tgt_len=input_shape[-1]).to(
133
- attention_mask_2d.device
134
- )
135
-
136
- if causal_4d_mask is not None:
137
- expanded_attn_mask = causal_4d_mask.masked_fill(expanded_attn_mask.bool(), torch.finfo(dtype).min)
138
-
139
- # expanded_attn_mask + causal_4d_mask can cause some overflow
140
- expanded_4d_mask = expanded_attn_mask
141
-
142
- return expanded_4d_mask
143
-
144
- @staticmethod
145
- def _make_causal_mask(
146
- input_ids_shape: torch.Size,
147
- dtype: torch.dtype,
148
- device: torch.device,
149
- past_key_values_length: int = 0,
150
- sliding_window: Optional[int] = None,
151
- ):
152
- """
153
- Make causal mask used for bi-directional self-attention.
154
- """
155
- bsz, tgt_len = input_ids_shape
156
- mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device)
157
- mask_cond = torch.arange(mask.size(-1), device=device)
158
- mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
159
-
160
- mask = mask.to(dtype)
161
-
162
- if past_key_values_length > 0:
163
- mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
164
-
165
- # add lower triangular sliding window mask if necessary
166
- if sliding_window is not None:
167
- diagonal = past_key_values_length - sliding_window + 1
168
-
169
- context_mask = 1 - torch.triu(torch.ones_like(mask, dtype=torch.int), diagonal=diagonal)
170
- mask.masked_fill_(context_mask.bool(), torch.finfo(dtype).min)
171
-
172
- return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
173
-
174
- @staticmethod
175
- def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
176
- """
177
- Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
178
- """
179
- bsz, src_len = mask.size()
180
- tgt_len = tgt_len if tgt_len is not None else src_len
181
-
182
- expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
183
-
184
- inverted_mask = 1.0 - expanded_mask
185
-
186
- return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)
187
-
188
- @staticmethod
189
- def _unmask_unattended(
190
- expanded_mask: torch.Tensor, attention_mask: torch.Tensor, unmasked_value: Union[bool, float]
191
- ):
192
- # fmt: off
193
- """
194
- Attend to all tokens in masked rows from the expanded attention mask, for example the relevant first rows when
195
- using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
196
- Details: https://github.com/pytorch/pytorch/issues/110213
197
-
198
- `expanded_mask` is [bsz, num_masks, tgt_seq_len, src_seq_len] or [bsz, tgt_seq_len, src_seq_len].
199
- `attention_mask` is [bsz, src_seq_len].
200
-
201
- The dimension num_masks of `expanded_mask` is most often 1, but it can also be the number of heads in the case of alibi attention bias.
202
-
203
- For example, if `attention_mask` is
204
- ```
205
- [[0, 0, 1],
206
- [1, 1, 1],
207
- [0, 1, 1]]
208
- ```
209
- and `expanded_mask` is (e.g. here left-padding case)
210
- ```
211
- [[[[0, 0, 0],
212
- [0, 0, 0],
213
- [0, 0, 1]]],
214
- [[[1, 0, 0],
215
- [1, 1, 0],
216
- [1, 1, 1]]],
217
- [[[0, 0, 0],
218
- [0, 1, 0],
219
- [0, 1, 1]]]]
220
- ```
221
- then the modified `expanded_mask` will be
222
- ```
223
- [[[[1, 1, 1], <-- modified
224
- [1, 1, 1], <-- modified
225
- [0, 0, 1]]],
226
- [[[1, 0, 0],
227
- [1, 1, 0],
228
- [1, 1, 1]]],
229
- [[[1, 1, 1], <-- modified
230
- [0, 1, 0],
231
- [0, 1, 1]]]]
232
- ```
233
- """
234
- # fmt: on
235
-
236
- # Get the index of the first non-zero value for every sample in the batch.
237
- # In the above example, indices = [[2], [0], [1]]]
238
- tmp = torch.arange(attention_mask.shape[1], 0, -1)
239
- indices = torch.argmax(attention_mask.cpu() * tmp, 1, keepdim=True)
240
-
241
- # Find the batch indexes that have unattended tokens on the leftmost side (e.g. [0, 0, 1, 1, 1]), for which the first rows of the
242
- # expanded mask will be completely unattended.
243
- left_masked_rows = torch.where(indices > 0)[0]
244
-
245
- if left_masked_rows.shape[0] == 0:
246
- return expanded_mask
247
- indices = indices[left_masked_rows]
248
-
249
- max_len = torch.max(indices)
250
- range_tensor = torch.arange(max_len).unsqueeze(0)
251
- range_tensor = range_tensor.repeat(indices.size(0), 1)
252
-
253
- # Avoid unmasking tokens at relevant target positions (on the row axis), by rather unmasking possibly several times the first row that should always be unmasked as we filtered out the batch above.
254
- range_tensor[range_tensor >= indices] = 0
255
-
256
- # TODO: we may drop support for 3D attention mask as the refactor from Patrick maybe dropped this case
257
- if expanded_mask.dim() == 4:
258
- num_masks = expanded_mask.shape[1]
259
- if num_masks == 1:
260
- # Broadcast [left_masked_rows, 1], [left_masked_rows, max_len]
261
- mask_slice = (left_masked_rows[:, None], 0, range_tensor)
262
- else:
263
- # Broadcast [left_masked_rows, 1, 1], [1, num_masks, 1], [left_masked_rows, 1, max_len]
264
- mask_slice = (
265
- left_masked_rows[:, None, None],
266
- torch.arange(num_masks)[None, :, None],
267
- range_tensor[:, None, :],
268
- )
269
- else:
270
- # Broadcast [left_masked_rows, 1], [left_masked_rows, max_len]
271
- mask_slice = (left_masked_rows[:, None], range_tensor)
272
-
273
- expanded_mask[mask_slice] = unmasked_value
274
-
275
- return expanded_mask
276
-
277
-
278
- def _prepare_4d_causal_attention_mask(
279
- attention_mask: Optional[torch.Tensor],
280
- input_shape: Union[torch.Size, Tuple, List],
281
- inputs_embeds: torch.Tensor,
282
- past_key_values_length: int,
283
- sliding_window: Optional[int] = None,
284
- ):
285
- """
286
- Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
287
- `(batch_size, key_value_length)`
288
-
289
- Args:
290
- attention_mask (`torch.Tensor` or `None`):
291
- A 2D attention mask of shape `(batch_size, key_value_length)`
292
- input_shape (`tuple(int)` or `list(int)` or `torch.Size`):
293
- The input shape should be a tuple that defines `(batch_size, query_length)`.
294
- inputs_embeds (`torch.Tensor`):
295
- The embedded inputs as a torch Tensor.
296
- past_key_values_length (`int`):
297
- The length of the key value cache.
298
- sliding_window (`int`, *optional*):
299
- If the model uses windowed attention, a sliding window should be passed.
300
- """
301
- attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window)
302
-
303
- key_value_length = input_shape[-1] + past_key_values_length
304
-
305
- # 4d mask is passed through the layers
306
- if attention_mask is not None and len(attention_mask.shape) == 2:
307
- attention_mask = attn_mask_converter.to_4d(
308
- attention_mask, input_shape[-1], key_value_length=key_value_length, dtype=inputs_embeds.dtype
309
- )
310
- elif attention_mask is not None and len(attention_mask.shape) == 4:
311
- expected_shape = (input_shape[0], 1, input_shape[1], key_value_length)
312
- if tuple(attention_mask.shape) != expected_shape:
313
- raise ValueError(
314
- f"Incorrect 4D attention_mask shape: {tuple(attention_mask.shape)}; expected: {expected_shape}."
315
- )
316
- else:
317
- # if the 4D mask has correct shape - invert it and fill with negative infinity
318
- inverted_mask = 1.0 - attention_mask
319
- attention_mask = inverted_mask.masked_fill(
320
- inverted_mask.to(torch.bool), torch.finfo(inputs_embeds.dtype).min
321
- )
322
- else:
323
- attention_mask = attn_mask_converter.to_causal_4d(
324
- input_shape[0], input_shape[-1], key_value_length, dtype=inputs_embeds.dtype, device=inputs_embeds.device
325
- )
326
-
327
- return attention_mask
328
-
329
-
330
- # Adapted from _prepare_4d_causal_attention_mask
331
- def _prepare_4d_causal_attention_mask_for_sdpa(
332
- attention_mask: Optional[torch.Tensor],
333
- input_shape: Union[torch.Size, Tuple, List],
334
- inputs_embeds: torch.Tensor,
335
- past_key_values_length: int,
336
- sliding_window: Optional[int] = None,
337
- ):
338
- """
339
- Prepares the correct `attn_mask` argument to be used by `torch.nn.functional.scaled_dot_product_attention`.
340
-
341
- In case no token is masked in the `attention_mask` argument, we simply set it to `None` for the cases `query_length == 1` and
342
- `key_value_length == query_length`, and rely instead on SDPA `is_causal` argument to use causal/non-causal masks,
343
- allowing to dispatch to the flash attention kernel (that can otherwise not be used if a custom `attn_mask` is passed).
344
- """
345
- attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window)
346
-
347
- key_value_length = input_shape[-1] + past_key_values_length
348
- batch_size, query_length = input_shape
349
-
350
- # torch.jit.trace, symbolic_trace and torchdynamo with fullgraph=True are unable to capture the controlflow `is_causal=attention_mask is None and q_len > 1`
351
- # used as an SDPA argument. We keep compatibility with these tracing tools by always using SDPA's `attn_mask` argument in case we are tracing.
352
- # TODO: Fix this as well when using torchdynamo with fullgraph=True.
353
- is_tracing = torch.jit.is_tracing() or isinstance(inputs_embeds, torch.fx.Proxy)
354
-
355
- if attention_mask is not None:
356
- # 4d mask is passed through
357
- if len(attention_mask.shape) == 4:
358
- expected_shape = (input_shape[0], 1, input_shape[1], key_value_length)
359
- if tuple(attention_mask.shape) != expected_shape:
360
- raise ValueError(
361
- f"Incorrect 4D attention_mask shape: {tuple(attention_mask.shape)}; expected: {expected_shape}."
362
- )
363
- else:
364
- # if the 4D mask has correct shape - invert it and fill with negative infinity
365
- inverted_mask = 1.0 - attention_mask.to(inputs_embeds.dtype)
366
- attention_mask = inverted_mask.masked_fill(
367
- inverted_mask.to(torch.bool), torch.finfo(inputs_embeds.dtype).min
368
- )
369
- return attention_mask
370
-
371
- elif not is_tracing and torch.all(attention_mask == 1):
372
- if query_length == 1:
373
- # For query_length == 1, causal attention and bi-directional attention are the same.
374
- attention_mask = None
375
- elif key_value_length == query_length:
376
- attention_mask = None
377
- else:
378
- # Unfortunately, for query_length > 1 and key_value_length != query_length, we cannot generally ignore the attention mask, as SDPA causal mask generation
379
- # may be wrong. We will set `is_causal=False` in SDPA and rely on Transformers attention_mask instead, hence not setting it to None here.
380
- # Reference: https://github.com/pytorch/pytorch/issues/108108
381
- pass
382
- elif query_length > 1 and key_value_length != query_length:
383
- # See the comment above (https://github.com/pytorch/pytorch/issues/108108).
384
- # Ugly: we set it to True here to dispatch in the following controlflow to `to_causal_4d`.
385
- attention_mask = True
386
- elif is_tracing:
387
- raise ValueError(
388
- 'Attention using SDPA can not be traced with torch.jit.trace when no attention_mask is provided. To solve this issue, please either load your model with the argument `attn_implementation="eager"` or pass an attention_mask input when tracing the model.'
389
- )
390
-
391
- if attention_mask is None:
392
- expanded_4d_mask = None
393
- elif attention_mask is True:
394
- expanded_4d_mask = attn_mask_converter.to_causal_4d(
395
- input_shape[0], input_shape[-1], key_value_length, dtype=inputs_embeds.dtype, device=inputs_embeds.device
396
- )
397
- else:
398
- expanded_4d_mask = attn_mask_converter.to_4d(
399
- attention_mask,
400
- input_shape[-1],
401
- dtype=inputs_embeds.dtype,
402
- key_value_length=key_value_length,
403
- )
404
-
405
- # From PyTorch 2.1 onwards, F.scaled_dot_product_attention with the memory-efficient attention backend
406
- # produces nans if sequences are completely unattended in the attention mask. Details: https://github.com/pytorch/pytorch/issues/110213
407
- #
408
- # This fix is not applied in case we are tracing with torch.jit.trace or symbolic_trace, as _unmask_unattended has a data-dependent
409
- # controlflow that can not be captured properly.
410
- # TODO: _unmask_unattended does not work either with torch.compile when using fullgraph=True. We should find a way to detect this case.
411
- if query_length > 1 and not is_tracing:
412
- expanded_4d_mask = AttentionMaskConverter._unmask_unattended(
413
- expanded_4d_mask, attention_mask, unmasked_value=0.0
414
- )
415
-
416
- return expanded_4d_mask
417
-
418
-
419
- def _prepare_4d_attention_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
420
- """
421
- Creates a non-causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
422
- `(batch_size, key_value_length)`
423
-
424
- Args:
425
- mask (`torch.Tensor` or `None`):
426
- A 2D attention mask of shape `(batch_size, key_value_length)`
427
- dtype (`torch.dtype`):
428
- The torch dtype the created mask shall have.
429
- tgt_len (`int`):
430
- The target length or query length the created mask shall have.
431
- """
432
- return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)
433
-
434
-
435
- def _prepare_4d_attention_mask_for_sdpa(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
436
- """
437
- Creates a non-causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
438
- `(batch_size, key_value_length)`
439
-
440
- Args:
441
- mask (`torch.Tensor` or `None`):
442
- A 2D attention mask of shape `(batch_size, key_value_length)`
443
- dtype (`torch.dtype`):
444
- The torch dtype the created mask shall have.
445
- tgt_len (`int`):
446
- The target length or query length the created mask shall have.
447
- """
448
- batch_size, key_value_length = mask.shape
449
- tgt_len = tgt_len if tgt_len is not None else key_value_length
450
-
451
- # torch.jit.trace and torchdynamo with fullgraph=True are unable to capture the controlflow `is_causal=attention_mask is None and q_len > 1`
452
- # used as an SDPA argument. We keep compatibility with these tracing tools by always using SDPA's `attn_mask` argument in case we are tracing.
453
- # TODO: Fix this as well when using torchdynamo with fullgraph=True.
454
- is_tracing = torch.jit.is_tracing()
455
-
456
- if torch.all(mask == 1):
457
- if is_tracing:
458
- pass
459
- elif tgt_len == 1:
460
- # For query_length == 1, causal attention and bi-directional attention are the same.
461
- return None
462
- elif key_value_length == tgt_len:
463
- return None
464
- else:
465
- # Unfortunately, for query_length > 1 and key_value_length != query_length, we can not generally ignore the attention mask, as SDPA causal mask generation
466
- # may be wrong. We will set is_causal=False in SDPA and rely on Transformers attention_mask instead, hence not setting it to None here.
467
- # Reference: https://github.com/pytorch/pytorch/issues/108108
468
- return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)
469
- else:
470
- return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)
471
-
472
-
473
- def _create_4d_causal_attention_mask(
474
- input_shape: Union[torch.Size, Tuple, List],
475
- dtype: torch.dtype,
476
- device: torch.device,
477
- past_key_values_length: int = 0,
478
- sliding_window: Optional[int] = None,
479
- ) -> Optional[torch.Tensor]:
480
- """
481
- Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)`
482
-
483
- Args:
484
- input_shape (`tuple(int)` or `list(int)` or `torch.Size`):
485
- The input shape should be a tuple that defines `(batch_size, query_length)`.
486
- dtype (`torch.dtype`):
487
- The torch dtype the created mask shall have.
488
- device (`int`):
489
- The torch device the created mask shall have.
490
- sliding_window (`int`, *optional*):
491
- If the model uses windowed attention, a sliding window should be passed.
492
- """
493
- attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window)
494
-
495
- key_value_length = past_key_values_length + input_shape[-1]
496
- attention_mask = attn_mask_converter.to_causal_4d(
497
- input_shape[0], input_shape[-1], key_value_length, dtype=dtype, device=device
498
- )
499
-
500
- return attention_mask