teowu commited on
Commit
728ace9
1 Parent(s): 15f7bb4

Update modeling_attn_mask_utils.py

Browse files
Files changed (1) hide show
  1. modeling_attn_mask_utils.py +88 -1
modeling_attn_mask_utils.py CHANGED
@@ -244,4 +244,91 @@ def _create_4d_causal_attention_mask(
244
  input_shape[0], input_shape[-1], key_value_length, dtype=dtype, device=device
245
  )
246
 
247
- return attention_mask
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
244
  input_shape[0], input_shape[-1], key_value_length, dtype=dtype, device=device
245
  )
246
 
247
+ return attention_mask
248
+
249
+
250
+ # Adapted from _prepare_4d_causal_attention_mask
251
+ def _prepare_4d_causal_attention_mask_for_sdpa(
252
+ attention_mask: Optional[torch.Tensor],
253
+ input_shape: Union[torch.Size, Tuple, List],
254
+ inputs_embeds: torch.Tensor,
255
+ past_key_values_length: int,
256
+ sliding_window: Optional[int] = None,
257
+ ):
258
+ """
259
+ Prepares the correct `attn_mask` argument to be used by `torch.nn.functional.scaled_dot_product_attention`.
260
+
261
+ In case no token is masked in the `attention_mask` argument, we simply set it to `None` for the cases `query_length == 1` and
262
+ `key_value_length == query_length`, and rely instead on SDPA `is_causal` argument to use causal/non-causal masks,
263
+ allowing to dispatch to the flash attention kernel (that can otherwise not be used if a custom `attn_mask` is passed).
264
+ """
265
+ attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window)
266
+
267
+ key_value_length = input_shape[-1] + past_key_values_length
268
+ batch_size, query_length = input_shape
269
+
270
+ # torch.jit.trace and torchdynamo with fullgraph=True are unable to capture the controlflow `is_causal=attention_mask is None and q_len > 1`
271
+ # 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.
272
+ # TODO: Fix this as well when using torchdynamo with fullgraph=True.
273
+ is_tracing = torch.jit.is_tracing()
274
+
275
+ if attention_mask is not None:
276
+ # 4d mask is passed through
277
+ if len(attention_mask.shape) == 4:
278
+ expected_shape = (input_shape[0], 1, input_shape[1], key_value_length)
279
+ if tuple(attention_mask.shape) != expected_shape:
280
+ raise ValueError(
281
+ f"Incorrect 4D attention_mask shape: {tuple(attention_mask.shape)}; expected: {expected_shape}."
282
+ )
283
+ else:
284
+ # if the 4D mask has correct shape - invert it and fill with negative infinity
285
+ inverted_mask = 1.0 - attention_mask.to(inputs_embeds.dtype)
286
+ attention_mask = inverted_mask.masked_fill(
287
+ inverted_mask.to(torch.bool), torch.finfo(inputs_embeds.dtype).min
288
+ )
289
+ return attention_mask
290
+
291
+ elif torch.all(attention_mask == 1):
292
+ if is_tracing:
293
+ pass
294
+ elif query_length == 1:
295
+ # For query_length == 1, causal attention and bi-directional attention are the same.
296
+ attention_mask = None
297
+ elif key_value_length == query_length:
298
+ attention_mask = None
299
+ else:
300
+ # Unfortunately, for query_length > 1 and key_value_length != query_length, we cannot generally ignore the attention mask, as SDPA causal mask generation
301
+ # 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.
302
+ # Reference: https://github.com/pytorch/pytorch/issues/108108
303
+ pass
304
+ elif query_length > 1 and key_value_length != query_length:
305
+ # See the comment above (https://github.com/pytorch/pytorch/issues/108108).
306
+ # Ugly: we set it to True here to dispatch in the following controlflow to `to_causal_4d`.
307
+ attention_mask = True
308
+ elif is_tracing:
309
+ raise ValueError(
310
+ '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.'
311
+ )
312
+
313
+ if attention_mask is None:
314
+ expanded_4d_mask = None
315
+ elif attention_mask is True:
316
+ expanded_4d_mask = attn_mask_converter.to_causal_4d(
317
+ input_shape[0], input_shape[-1], key_value_length, dtype=inputs_embeds.dtype, device=inputs_embeds.device
318
+ )
319
+ else:
320
+ expanded_4d_mask = attn_mask_converter.to_4d(
321
+ attention_mask,
322
+ input_shape[-1],
323
+ dtype=inputs_embeds.dtype,
324
+ key_value_length=key_value_length,
325
+ )
326
+
327
+ # From PyTorch 2.1 onwards, F.scaled_dot_product_attention with the memory-efficient attention backend
328
+ # produces nans if sequences are completely unattended in the attention mask. Details: https://github.com/pytorch/pytorch/issues/110213
329
+ if query_length > 1:
330
+ expanded_4d_mask = AttentionMaskConverter._unmask_unattended(
331
+ expanded_4d_mask, attention_mask, unmasked_value=0.0
332
+ )
333
+
334
+ return expanded_4d_mask