Skip to content

Phi3: fix attn for sliding window #33586

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zucchini-nlp merged 10 commits into from
Oct 10, 2024
Merged

Phi3: fix attn for sliding window #33586

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
156a85b
fix phi3 attn fir sliding window
zucchini-nlp Sep 19, 2024
2e33d0d
fix tests
zucchini-nlp Sep 19, 2024
0faeaf3
merge main
zucchini-nlp Sep 26, 2024
52fb350
address most comment
zucchini-nlp Sep 26, 2024
2351ab2
style
zucchini-nlp Sep 26, 2024
5b27f5f
merge main
zucchini-nlp Oct 8, 2024
a19c7f5
update after rebase
zucchini-nlp Oct 8, 2024
c174f25
add more models
zucchini-nlp Oct 9, 2024
21ede7f
merge main
zucchini-nlp Oct 9, 2024
25b5852
fix tests
zucchini-nlp Oct 9, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion src/transformers/cache_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1244,7 +1244,6 @@ def __init__(
max_batch_size: Optional[int] = None,
layer_device_map: Optional[Dict[int, Union[str, torch.device, int]]] = None,
) -> None:
super().__init__()
if not hasattr(config, "sliding_window") or config.sliding_window is None:
raise ValueError(
"Setting `cache_implementation` to 'sliding_window' requires the model config supporting "
Expand Down
151 changes: 118 additions & 33 deletions src/transformers/models/mistral/modeling_mistral.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,72 @@
_CONFIG_FOR_DOC = "MistralConfig"


# Copied from transformers.models.phi3.modeling_phi3._prepare_4d_causal_attention_mask_with_cache_position with Phi3->Mistral
Comment thread
zucchini-nlp marked this conversation as resolved.
Outdated
def _prepare_4d_causal_attention_mask_with_cache_position(

@gante gante Oct 4, 2024

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Because of #33677, this function is part of the model class -- I think you will have to move the diff there, otherwise tests may fail on main

(see the diff in that PR for Llama, it should be similar to the changes you need to do here)

attention_mask: torch.Tensor,
sequence_length: int,
target_length: int,
dtype: torch.dtype,
device: torch.device,
min_dtype: float,
cache_position: torch.Tensor,
batch_size: int,
config: MistralConfig,
past_key_values: Cache,
):
"""
Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
`(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.

Args:
attention_mask (`torch.Tensor`):
A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`.
sequence_length (`int`):
The sequence length being processed.
target_length (`int`):
The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet.
dtype (`torch.dtype`):
The dtype to use for the 4D attention mask.
device (`torch.device`):
The device to plcae the 4D attention mask on.
min_dtype (`float`):
The minimum value representable with the dtype `dtype`.
cache_position (`torch.Tensor`):
Indices depicting the position of the input sequence tokens in the sequence.
batch_size (`torch.Tensor`):
Batch size.
config (`MistralConfig`):
The model's configuration class
past_key_values (`Cache`):
The cache class that is being used currently to generate
"""
if attention_mask is not None and attention_mask.dim() == 4:
# In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.
causal_mask = attention_mask
else:
causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device)
exclude_mask = torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
if config.sliding_window is not None:
if not isinstance(past_key_values, SlidingWindowCache) or sequence_length > config.sliding_window:
exclude_mask.bitwise_or_(
torch.arange(target_length, device=device)
<= (cache_position.reshape(-1, 1) - config.sliding_window)
)
causal_mask *= exclude_mask
causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
if attention_mask is not None:
causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit
if attention_mask.shape[-1] > target_length:
attention_mask = attention_mask[:, :target_length]
mask_length = attention_mask.shape[-1]
padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :]
padding_mask = padding_mask == 0
causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
padding_mask, min_dtype
)
return causal_mask


# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Mistral
class MistralRMSNorm(nn.Module):
def __init__(self, hidden_size, eps=1e-6):
Expand Down Expand Up @@ -853,7 +919,7 @@ def _update_causal_mask(
use_cache: bool,
output_attentions: bool,
):
if self._attn_implementation == "flash_attention_2":
if self.config._attn_implementation == "flash_attention_2":
if attention_mask is not None and use_cache:
is_padding_right = attention_mask[:, -1].sum().item() != input_tensor.size()[0]
if is_padding_right:
Expand All @@ -869,12 +935,11 @@ def _update_causal_mask(
# For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in
# order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail
# to infer the attention mask.

# cache_position must be valid here no matter which cache we use
past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
using_static_cache = isinstance(past_key_values, StaticCache)
using_sliding_window_cache = isinstance(past_key_values, SlidingWindowCache)

# When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward
if (
self.config._attn_implementation == "sdpa"
and not (using_static_cache or using_sliding_window_cache)
Expand All @@ -892,44 +957,31 @@ def _update_causal_mask(
dtype, device = input_tensor.dtype, input_tensor.device
min_dtype = torch.finfo(dtype).min
sequence_length = input_tensor.shape[1]
# SlidingWindowCache
if using_sliding_window_cache:
target_length = max(sequence_length, self.config.sliding_window)
# StaticCache
sliding_window = self.config.sliding_window if self.config.sliding_window is not None else sequence_length
target_length = max(sequence_length, sliding_window)
elif using_static_cache:
target_length = past_key_values.get_max_length()
# DynamicCache or no cache
else:
target_length = (
attention_mask.shape[-1]
if isinstance(attention_mask, torch.Tensor)
else past_seen_tokens + sequence_length + 1
)

if attention_mask is not None and attention_mask.dim() == 4:
causal_mask = attention_mask
else:
causal_mask = torch.full(
(sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device
)
exclude_mask = torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
if self.config.sliding_window is not None:
if not using_sliding_window_cache or sequence_length > self.config.sliding_window:
exclude_mask.bitwise_or_(
torch.arange(target_length, device=device)
<= (cache_position.reshape(-1, 1) - self.config.sliding_window)
)
causal_mask *= exclude_mask
causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1)
if attention_mask is not None:
causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit
if attention_mask.dim() == 2:
mask_length = attention_mask.shape[-1]
padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :]
padding_mask = padding_mask == 0
causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
padding_mask, min_dtype
)
# In case the provided `attention` mask is 2D, we generate a causal mask here (4D).
causal_mask = _prepare_4d_causal_attention_mask_with_cache_position(
attention_mask,
sequence_length=sequence_length,
target_length=target_length,
dtype=dtype,
device=device,
min_dtype=min_dtype,
cache_position=cache_position,
batch_size=input_tensor.shape[0],
config=self.config,
past_key_values=past_key_values,
)

if (
self.config._attn_implementation == "sdpa"
Expand Down Expand Up @@ -1112,9 +1164,42 @@ def prepare_inputs_for_generation(

# if `inputs_embeds` are passed, we only want to use them in the 1st generation step
if inputs_embeds is not None and cache_position[0] == 0:
model_inputs = {"inputs_embeds": inputs_embeds}
model_inputs = {"inputs_embeds": inputs_embeds, "input_ids": None}
Comment thread
zucchini-nlp marked this conversation as resolved.
else:
model_inputs = {"input_ids": input_ids.contiguous()} # `contiguous()` needed for compilation use cases
# `contiguous()` needed for compilation use cases
model_inputs = {"input_ids": input_ids.contiguous(), "inputs_embeds": None}

if isinstance(past_key_values, StaticCache) and attention_mask.ndim == 2:
if model_inputs["inputs_embeds"] is not None:
batch_size, sequence_length, _ = model_inputs["inputs_embeds"].shape
device = model_inputs["inputs_embeds"].device
else:
batch_size, sequence_length = model_inputs["input_ids"].shape
device = model_inputs["input_ids"].device

dtype = self.lm_head.weight.dtype
min_dtype = torch.finfo(dtype).min

if isinstance(past_key_values, SlidingWindowCache):
sliding_window = (
self.config.sliding_window if self.config.sliding_window is not None else sequence_length
)
target_length = max(sequence_length, sliding_window)
else:
target_length = past_key_values.get_max_length()

@ArthurZucker ArthurZucker Oct 3, 2024

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't understand why we need to do this: if there is a sliding WindowCache it was init from the config and thus has a correct sliding_window.

Then, SlidingWindowCache. get_max_length should take sequence_length as input to return the max and avoid having these checks here WDYT?

@zucchini-nlp zucchini-nlp Oct 3, 2024

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, it should but the PR for getting max_length on slidingWindow cache is on its way and not merged yet.

So, to update you on our discussions with @gante which are currently in different PR comments: some cache classes now do not have a max_length (e.g. SlidingWindow). As commented in the code, sliding window technically has no max length and goes on a rolling basis. But in transformers what we want to check is the "maximum capacity of cache instance", independently of how cache handles new tokens going beyond that capacity.

So, in a different PR I changed naming to get_max_cache_shape which is more straightforward and added get_max_cache_shape for Sliding Window cache. We'll do a simple deprecation cycle, as we did for static cache's "max_batch_size". Until the linked PR is merged, I am copying this piece of code from mistral and using it in phi3. I have it noted and will handle it depending which one gets merged first :)

@ArthurZucker ArthurZucker Oct 4, 2024

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

alright sounds good. I just don't want us to add too much complexity to the code! 🤗


attention_mask = _prepare_4d_causal_attention_mask_with_cache_position(
attention_mask,
sequence_length=sequence_length,
target_length=target_length,
dtype=dtype,
device=device,
min_dtype=min_dtype,
cache_position=cache_position,
batch_size=batch_size,
config=self.config,
past_key_values=past_key_values,
)

if num_logits_to_keep is not None:
model_inputs["num_logits_to_keep"] = num_logits_to_keep
Expand Down
51 changes: 41 additions & 10 deletions src/transformers/models/phi3/modeling_phi3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACT2FN
from ...cache_utils import Cache, DynamicCache, StaticCache
from ...cache_utils import Cache, DynamicCache, SlidingWindowCache, StaticCache
from ...modeling_attn_mask_utils import AttentionMaskConverter
from ...modeling_outputs import (
BaseModelOutputWithPast,
Expand Down Expand Up @@ -56,7 +56,6 @@
_CONFIG_FOR_DOC = "Phi3Config"


# Copied from transformers.models.llama.modeling_llama._prepare_4d_causal_attention_mask_with_cache_position
def _prepare_4d_causal_attention_mask_with_cache_position(
attention_mask: torch.Tensor,
sequence_length: int,
Expand All @@ -66,6 +65,8 @@ def _prepare_4d_causal_attention_mask_with_cache_position(
min_dtype: float,
cache_position: torch.Tensor,
batch_size: int,
config: Phi3Config,
past_key_values: Cache,
):
"""
Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
Expand All @@ -88,25 +89,35 @@ def _prepare_4d_causal_attention_mask_with_cache_position(
Indices depicting the position of the input sequence tokens in the sequence.
batch_size (`torch.Tensor`):
Batch size.
config (`Phi3Config`):
The model's configuration class
past_key_values (`Cache`):
The cache class that is being used currently to generate
"""
if attention_mask is not None and attention_mask.dim() == 4:
# In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.
causal_mask = attention_mask
else:
causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device)
if sequence_length != 1:
causal_mask = torch.triu(causal_mask, diagonal=1)
causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
exclude_mask = torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
Comment thread
zucchini-nlp marked this conversation as resolved.
Outdated
if config.sliding_window is not None:
if not isinstance(past_key_values, SlidingWindowCache) or sequence_length > config.sliding_window:
exclude_mask.bitwise_or_(
torch.arange(target_length, device=device)
<= (cache_position.reshape(-1, 1) - config.sliding_window)
Comment thread
zucchini-nlp marked this conversation as resolved.
Outdated
)
causal_mask *= exclude_mask
causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
if attention_mask is not None:
causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit
if attention_mask.shape[-1] > target_length:
attention_mask = attention_mask[:, :target_length]
Comment thread
zucchini-nlp marked this conversation as resolved.
Outdated
mask_length = attention_mask.shape[-1]
padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :]
padding_mask = padding_mask == 0
causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
padding_mask, min_dtype
)

return causal_mask


Expand Down Expand Up @@ -1086,7 +1097,6 @@ def forward(
attentions=all_self_attns,
)

# Copied from transformers.models.llama.modeling_llama.LlamaModel._update_causal_mask
def _update_causal_mask(
self,
attention_mask: torch.Tensor,
Expand All @@ -1105,21 +1115,30 @@ def _update_causal_mask(
# to infer the attention mask.
past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
using_static_cache = isinstance(past_key_values, StaticCache)
using_sliding_window_cache = isinstance(past_key_values, SlidingWindowCache)

# When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward
if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions:
if (
self.config._attn_implementation == "sdpa"
and not (using_static_cache or using_sliding_window_cache)
and not output_attentions
):
if AttentionMaskConverter._ignore_causal_mask_sdpa(
attention_mask,
inputs_embeds=input_tensor,
past_key_values_length=past_seen_tokens,
sliding_window=self.config.sliding_window,
is_training=self.training,
):
return None

dtype, device = input_tensor.dtype, input_tensor.device
min_dtype = torch.finfo(dtype).min
sequence_length = input_tensor.shape[1]
if using_static_cache:
if using_sliding_window_cache:

@gante gante Sep 19, 2024

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we now have a significant code block to determine target_length, but target_length is not used directly in _update_causal_mask

suggestion to avoid this disconnection:

  1. make SlidingWindowCache store the config.sliding_window it receives at init time
  2. move the logic that computes target_length inside _prepare_4d_causal_attention_mask_with_cache_position
  3. _prepare_4d_causal_attention_mask_with_cache_position no longer receives target_length nor config, as they can be retrieved from past_key_values

@zucchini-nlp zucchini-nlp Sep 19, 2024

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For the target length, maybe we merge #32421 first where the max length should be easily accessible through cache no matter cache type. We earlier discussed returning max cache shape, and sliding window has a max capacity for cache class

And then yeah, we can move that part to _prepare_4d_causal_attention_mask_with_cache_position but I think we'd better move it in all classes for general consistency

@gante gante Sep 19, 2024

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed with the plan -- shall we leave a TODO for us in this PR, linking to your comment?

@zucchini-nlp zucchini-nlp Sep 26, 2024

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes, I'll add a TODO for us so we can make the required change in all models at once, to not mix different updates in one PR :)

sliding_window = self.config.sliding_window if self.config.sliding_window is not None else sequence_length
target_length = max(sequence_length, sliding_window)
elif using_static_cache:
target_length = past_key_values.get_max_length()
else:
target_length = (
Expand All @@ -1138,6 +1157,8 @@ def _update_causal_mask(
min_dtype=min_dtype,
cache_position=cache_position,
batch_size=input_tensor.shape[0],
config=self.config,
past_key_values=past_key_values,
)

if (
Expand Down Expand Up @@ -1364,15 +1385,25 @@ def prepare_inputs_for_generation(
dtype = self.lm_head.weight.dtype
min_dtype = torch.finfo(dtype).min

if isinstance(past_key_values, SlidingWindowCache):

@gante gante Sep 19, 2024

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(same comment as above here)

sliding_window = (
self.config.sliding_window if self.config.sliding_window is not None else sequence_length
)
target_length = max(sequence_length, sliding_window)
else:
target_length = past_key_values.get_max_length()

attention_mask = _prepare_4d_causal_attention_mask_with_cache_position(
attention_mask,
sequence_length=sequence_length,
target_length=past_key_values.get_max_length(),
target_length=target_length,
dtype=dtype,
device=device,
min_dtype=min_dtype,
cache_position=cache_position,
batch_size=batch_size,
config=self.config,
past_key_values=past_key_values,
)

if num_logits_to_keep is not None:
Expand Down
Loading