(query_layer, key_layer, value_layer, attention_mask,
attention_dropout=None, log_attention_weights=None, scaling_attention_score=True, **kwargs)
| 45 | return context_layer |
| 46 | |
| 47 | def attention_fn_default(query_layer, key_layer, value_layer, attention_mask, |
| 48 | attention_dropout=None, log_attention_weights=None, scaling_attention_score=True, **kwargs): |
| 49 | # expand head dim to query dim, if necessary |
| 50 | # only useful for multi-query attention |
| 51 | batch_size, num_query_heads = query_layer.shape[:2] # [b, np, s, hn] |
| 52 | num_kv_heads = key_layer.shape[1] # [b, np, s, hn] |
| 53 | key_layer = key_layer.unsqueeze(2).expand(-1, -1, num_query_heads//num_kv_heads, -1, -1).contiguous().view(batch_size, num_query_heads, *key_layer.shape[2:]) |
| 54 | value_layer = value_layer.unsqueeze(2).expand(-1, -1, num_query_heads//num_kv_heads, -1, -1).contiguous().view(batch_size, num_query_heads, *value_layer.shape[2:]) |
| 55 | |
| 56 | is_low_triangle = (attention_mask == torch.ones_like(attention_mask, dtype=torch.float).tril()).all() |
| 57 | is_full = (attention_mask is None) or (attention_mask > 0).all() |
| 58 | |
| 59 | if int(torch.__version__.split('.')[0]) >= 2 and scaling_attention_score and (is_full or is_low_triangle): |
| 60 | # Pytorch 2.0 attention uses very much memory if attention_mask is float, and has NaN bug if attention_mask is None. |
| 61 | dropout_p = 0. if attention_dropout is None or not attention_dropout.training else attention_dropout.p |
| 62 | if dropout_p > 0 and mpu.get_cuda_rng_tracker is not None: |
| 63 | context = mpu.get_cuda_rng_tracker().fork() |
| 64 | else: |
| 65 | context = contextlib.nullcontext() |
| 66 | with context: |
| 67 | attn_output = torch.nn.functional.scaled_dot_product_attention( |
| 68 | query_layer, key_layer, value_layer, |
| 69 | attn_mask=None, |
| 70 | dropout_p=dropout_p, |
| 71 | is_causal=not is_full |
| 72 | ) |
| 73 | return attn_output |
| 74 | else: |
| 75 | return standard_attention( |
| 76 | query_layer, key_layer, value_layer, attention_mask, |
| 77 | attention_dropout=attention_dropout, log_attention_weights=log_attention_weights, |
| 78 | scaling_attention_score=scaling_attention_score, **kwargs |
| 79 | ) |
| 80 | |
| 81 | def attention_forward_default(self, hidden_states, mask, **kw_args): |
| 82 | self = self.transformer.layers[kw_args['layer_id']].attention |
no test coverage detected