Attn added kv processor torch 2.0 block #3023
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
The documentation is not available anymore as the PR was closed or merged. |
williamberman
force-pushed
the
AttnAddedKVProcessor2_0
branch
2 times, most recently
from
26db440 to
40ec593
Compare
williamberman
requested review from
patrickvonplaten,
pcuenca,
sayakpaul and
yiyixuxu
williamberman
force-pushed
the
AttnAddedKVProcessor2_0
branch
from
40ec593 to
a38867e
Compare
williamberman
commented
Apr 9, 2023
Comment on lines
+324
to
+333
| if out_dim == 3: | ||
| if attention_mask.shape[0] < batch_size * head_size: | ||
| attention_mask = attention_mask.repeat_interleave(head_size, dim=0) | ||
| elif out_dim == 4: | ||
| attention_mask = attention_mask.unsqueeze(1) | ||
| attention_mask = attention_mask.repeat_interleave(head_size, dim=1) | ||
|
|
There was a problem hiding this comment.
@patrickvonplaten when using the torch built in attention and putting heads as the second dim, we need to make the attention mask also put the heads in the second dim. I'm not sure what the equivalent check for attention_mask.shape[0] < batch_size * head_size is. If we assume the input attention mask is always just the same batch size of the inputs, we don't have to do the check and I think this works. My understanding is that's what the original code was doing anyway since it just repeats by the head size regardless
patrickvonplaten
approved these changes
Apr 9, 2023
williamberman
force-pushed
the
AttnAddedKVProcessor2_0
branch
from
a38867e to
a480bc0
Compare
sayakpaul
reviewed
Apr 11, 2023
Comment on lines
113
to
116
| else: | ||
| raise ValueError( | ||
| f"unknown cross_attention_norm: {cross_attention_norm}. Should be None, 'layer_norm' or 'group_norm'" | ||
| ) |
williamberman
force-pushed
the
AttnAddedKVProcessor2_0
branch
from
a480bc0 to
fb485ad
Compare
sayakpaul
reviewed
Apr 11, 2023
| value = encoder_hidden_states_value_proj | ||
|
|
||
| # the output of sdp = (batch, num_heads, seq_len, head_dim) | ||
| # TODO: add support for attn.scale when we move to Torch 2.1 |
There was a problem hiding this comment.
Out of curiosity. Which scale is this one?
Comment on lines
342
to
349
| elif isinstance(self.norm_cross, nn.GroupNorm): | ||
| # Group norm norms along the channels dimension and expects | ||
| # input to be in the shape of (N, C, *). In this case, we want | ||
| # to norm along the hidden dimension, so we need to move | ||
| # (batch_size, sequence_length, hidden_size) -> | ||
| # (batch_size, hidden_size, sequence_length) | ||
| encoder_hidden_states = encoder_hidden_states.transpose(1, 2) | ||
| encoder_hidden_states = self.norm_cross(encoder_hidden_states) | ||
| encoder_hidden_states = encoder_hidden_states.transpose(1, 2) |
Comment on lines
486
to
494
| if not attn.only_cross_attention: | ||
| key = attn.to_k(hidden_states) | ||
| value = attn.to_v(hidden_states) | ||
| key = attn.head_to_batch_dim(key) | ||
| value = attn.head_to_batch_dim(value) | ||
| key = torch.cat([encoder_hidden_states_key_proj, key], dim=1) | ||
| value = torch.cat([encoder_hidden_states_value_proj, value], dim=1) | ||
| else: | ||
| key = encoder_hidden_states_key_proj | ||
| value = encoder_hidden_states_value_proj |
Comment on lines
1523
to
1569
| only_cross_attention=only_cross_attention, | ||
| cross_attention_norm=cross_attention_norm, | ||
| processor=processor, |
Comment on lines
+424
to
+429
| # Check is relaxed because there is not a torch 2.0 sliced attention added kv processor | ||
| expected_max_diff = 1e-2 | ||
|
|
||
| self._test_attention_slicing_forward_pass( | ||
| test_max_difference=test_max_difference, expected_max_diff=expected_max_diff | ||
| ) |
There was a problem hiding this comment.
Should this be conditioned on the torch version being used?
There was a problem hiding this comment.
Sorry should have noted, intention is to follow up and add the 2.0 sliced attention added kv processor
williamberman
force-pushed
the
AttnAddedKVProcessor2_0
branch
from
fb485ad to
ac3786a
Compare
williamberman
force-pushed
the
AttnAddedKVProcessor2_0
branch
from
ac3786a to
a65dd58
Compare
w4ffl35
pushed a commit
to w4ffl35/diffusers
that referenced
this pull request
Apr 14, 2023
add AttnAddedKVProcessor2_0 block
dg845
pushed a commit
to dg845/diffusers
that referenced
this pull request
May 6, 2023
add AttnAddedKVProcessor2_0 block
yoonseokjin
pushed a commit
to yoonseokjin/diffusers
that referenced
this pull request
Dec 25, 2023
add AttnAddedKVProcessor2_0 block
AmericanPresidentJimmyCarter
pushed a commit
to AmericanPresidentJimmyCarter/diffusers
that referenced
this pull request
Apr 26, 2024
add AttnAddedKVProcessor2_0 block
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
rebased on top of #3021 and #3011
200a8c7 and on is the main commit