Skip to content

Expand FA coverage to padding masks#291

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

Closed
ksivaman wants to merge 12 commits into from
Closed

Expand FA coverage to padding masks #291

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
9a28af3
initial changes [wip]
ksivaman Jun 21, 2023
7b01a27
Merge branch 'main' into expand_flash_attn_coverage
ksivaman Jun 21, 2023
a89e43c
Add padding mask support for FA
ksivaman Jun 21, 2023
ae104b3
Address review comments
ksivaman Jun 21, 2023
e38d18d
rm causal mask from tests and add padding
ksivaman Jun 22, 2023
b0212bd
fix some conflicts
ksivaman Jun 22, 2023
e3dfd8b
conflicts
ksivaman Jun 22, 2023
ad8fc31
Merge branch 'main' into expand_flash_attn_coverage
ksivaman Jun 22, 2023
e2b2ddf
Fixes
ksivaman Jun 22, 2023
8eb8d90
Merge branch 'main' into expand_flash_attn_coverage
ksivaman Jun 23, 2023
e0cd408
Merge branch 'main' into expand_flash_attn_coverage
ksivaman Jun 26, 2023
e181188
Merge branch 'main' into expand_flash_attn_coverage
ksivaman Jun 27, 2023
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
13 changes: 5 additions & 8 deletions tests/pytorch/test_numerics.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -401,11 +401,10 @@ def _test_e2e_selective_recompute(block, bs, dtype, config, recompute=False):
config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
).cuda()
te_inp_hidden_states.retain_grad()
te_inp_attn_mask = get_causal_attn_mask(config.seq_len)

te_out = block(
te_inp_hidden_states,
te_inp_attn_mask,
None,
checkpoint_core_attention=recompute,
)
loss = te_out.sum()
Expand DownExpand Up@@ -461,7 +460,6 @@ def _test_e2e_full_recompute(block, bs, dtype, config, recompute=False):
config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
).cuda()
te_inp_hidden_states.retain_grad()
te_inp_attn_mask = get_causal_attn_mask(config.seq_len)

if recompute:
te_out = te_checkpoint(
Expand All@@ -470,13 +468,13 @@ def _test_e2e_full_recompute(block, bs, dtype, config, recompute=False):
get_dummy_cuda_rng_tracker,
None, # tp_group
te_inp_hidden_states,
te_inp_attn_mask,
None,
checkpoint_core_attention=False,
)
else:
te_out = block(
te_inp_hidden_states,
te_inp_attn_mask,
None,
checkpoint_core_attention=False,
)
loss = te_out.sum()
Expand DownExpand Up@@ -556,14 +554,13 @@ def _test_e2e_checkpointing(bs, dtype, config, checkpoint=False, steps=10, path=
config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
).cuda()
te_inp_hidden_states.retain_grad()
te_inp_attn_mask = get_causal_attn_mask(config.seq_len)

block = _test_e2e_checkpointing_get_model(config, dtype)

for _ in range(steps // 2):
te_out = block(
te_inp_hidden_states,
te_inp_attn_mask,
None,
)
loss = te_out.sum()
loss.backward()
Expand DownExpand Up@@ -594,7 +591,7 @@ def _test_e2e_checkpointing(bs, dtype, config, checkpoint=False, steps=10, path=
for _ in range(steps // 2):
te_out = block(
te_inp_hidden_states,
te_inp_attn_mask,
None,
)
loss = te_out.sum()
loss.backward()
Expand Down
91 changes: 39 additions & 52 deletions tests/pytorch/test_sanity.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -154,26 +154,14 @@ def _test_sanity_e2e_amp(block, bs, dtype, config, fp8_recipe, skip_wgrad):
config.seq_len, bs, config.hidden_size, dtype=torch.float32, requires_grad=True
).cuda()
te_inp_hidden_states.retain_grad()
te_inp_attn_mask = (
torch.rand(
(
1,
1,
config.seq_len,
config.seq_len,
)
)
.cuda()
.bool()
)

if skip_wgrad:
_disable_wgrads(block)

use_fp8 = fp8_recipe is not None
with torch.autocast(device_type="cuda", enabled=True, dtype=dtype):
with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
te_out = block(te_inp_hidden_states, te_inp_attn_mask)
te_out = block(te_inp_hidden_states)
loss = te_out.sum()

loss.backward()
Expand All@@ -190,18 +178,6 @@ def _test_sanity_e2e_gradient_accumulation_fusion(block, bs, dtype, config, fp8_
te_inp_hidden_states = torch.randn(
config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
).cuda()
te_inp_attn_mask = (
torch.rand(
(
1,
1,
config.seq_len,
config.seq_len,
)
)
.cuda()
.bool()
)

if skip_wgrad:
_disable_wgrads(block)
Expand All@@ -214,7 +190,7 @@ def _test_sanity_e2e_gradient_accumulation_fusion(block, bs, dtype, config, fp8_

use_fp8 = fp8_recipe is not None
with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
te_out = block(te_inp_hidden_states, te_inp_attn_mask)
te_out = block(te_inp_hidden_states)
loss = te_out.sum()
loss.backward()
torch.cuda.synchronize()
Expand All@@ -232,18 +208,29 @@ def _test_sanity_e2e(block, bs, dtype, config, fp8_recipe, skip_wgrad):
te_inp_hidden_states = torch.randn(
config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
).cuda()
te_inp_attn_mask = (
torch.rand(
(
1,
1,
config.seq_len,
config.seq_len,
)
)
.cuda()
.bool()
)

if skip_wgrad:
_disable_wgrads(block)

use_fp8 = fp8_recipe is not None
with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
te_out = block(te_inp_hidden_states)
loss = te_out.sum()
loss.backward()
torch.cuda.synchronize()


def _test_sanity_e2e_bert(block, bs, dtype, config, fp8_recipe, skip_wgrad):
te_inp_hidden_states = torch.randn(
config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
).cuda()

if dtype == torch.float32:
mask_shape = torch.Size([1, 1, config.seq_len, config.seq_len])
else:
mask_shape = torch.Size([config.seq_len, bs])

te_inp_attn_mask = torch.rand(mask_shape).cuda().bool()
Comment on lines +228 to +233

@timmoon10timmoon10Jun 21, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Why does the FP32 test require a 4D mask while other dtypes use a 2D mask?

@ksivamanksivamanJun 21, 2023

Copy link
Copy Markdown
MemberAuthor

Choose a reason for hiding this comment

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

The FP32 doesn't take the FA path and so uses the PyTorch torch softmax path for padding mask which expects this

@timmoon10timmoon10Jun 22, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Ah, so FA vs PyT is no longer an implementation detail since it expects a different mask format. This makes me think we need an option to explicitly enable or disable FA, and to error out instead of falling back to PyT.

@timmoon10timmoon10Jun 22, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Alternatively, we should expect the mask to be in the format for PyT, and then convert it to the FA format internally so it isn't visible to the user.

@ptrendxptrendxJun 23, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Agree with Tim. OR, accept both in both cases and generate the right one for the current implementation.

To be honest though, the current API for mask type is just bad. For example, what if you want padding AND causal? Then your only option is set the padding mask type and create the causal mask yourself (which would mean you can't do FA in the current implementation which is bad). There is also nothing stopping you really from having arbitrary padding mask (as in, with random elements zeroed out), which would break the assumptions you have in this PR. I think what we could do is introduce new names for the padding type (let's say "pad", "arbitrary" and "no_mask"), deprecate the old names and add causal as separate switch. Then we could say that pad only accepts the "mask" that is actually just list of sequence lengths, arbitrary is whatever you want (and will not go through FA) and causal could be switched irrespective of the mask type.


if skip_wgrad:
_disable_wgrads(block)
Expand All@@ -260,26 +247,24 @@ def _test_sanity_e2e_T5(block, bs, dtype, config, fp8_recipe, skip_wgrad):
te_inp_hidden_states = torch.randn(
config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
).cuda()
te_inp_attn_mask = (
torch.rand(
(
1,
1,
config.seq_len,
config.seq_len,
)
)
.cuda()
.bool()
)

if dtype == torch.float32:
mask_shape = torch.Size([1, 1, config.seq_len, config.seq_len])
else:
mask_shape = torch.Size([config.seq_len, bs])

enc_dec_attn_mask = torch.rand(mask_shape).cuda().bool()

if skip_wgrad:
_disable_wgrads(block)

use_fp8 = fp8_recipe is not None
with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
te_out = block(
te_inp_hidden_states, te_inp_attn_mask, encoder_output=te_inp_hidden_states
te_inp_hidden_states,
None,
encoder_output=te_inp_hidden_states,
enc_dec_attn_mask=enc_dec_attn_mask,
)
loss = te_out.sum()
loss.backward()
Expand DownExpand Up@@ -468,12 +453,14 @@ def test_sanity_bert(dtype, bs, fp8_recipe, model, skip_wgrad, zero_centered_gam
apply_residual_connection_post_layernorm=True,
output_layernorm=True,
zero_centered_gamma=zero_centered_gamma,
self_attn_mask_type="padding",

)
.to(dtype=dtype)
.cuda()
)

_test_sanity_e2e(block, bs, dtype, config, fp8_recipe, skip_wgrad)
_test_sanity_e2e_bert(block, bs, dtype, config, fp8_recipe, skip_wgrad)


@pytest.mark.parametrize("dtype", param_types)
Expand Down
61 changes: 41 additions & 20 deletions transformer_engine/pytorch/attention.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -44,6 +44,7 @@
checkpoint,
)
from transformer_engine.pytorch.export import is_in_onnx_export_mode
from transformer_engine.pytorch.jit import jit_fuser

_flash_attn_version = packaging.version.Version(version("flash-attn"))
_flash_attn_version_required = packaging.version.Version("1.0.6")
Expand All@@ -52,6 +53,19 @@
__all__ = ["DotProductAttention"]


@jit_fuser
def get_cu_seqlens(padding_mask: torch.Tensor) -> torch.Tensor:
"""
Given a padding mask of shape [seq_len, batch_size], returns an int32

@cyanguwacyanguwaJun 22, 2023

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 guess by seq_len you mean max_seqlen? Also, should the mask tensor be required to be a CUDA tensor?

@ksivamanksivamanJun 22, 2023

Copy link
Copy Markdown
MemberAuthor

Choose a reason for hiding this comment

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

Yeah it is the max_seqlen.

@ksivamanksivamanJun 22, 2023

Copy link
Copy Markdown
MemberAuthor

Choose a reason for hiding this comment

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

Yeah it should be cuda tensor since the inp is cuda as well.

tensor of shape [batch_size + 1,] containing the cumulative sequence
lengths of every sample in the batch.
"""
reduced_mask = padding_mask.sum(dim=0)
cu_seqlens = reduced_mask.cumsum(dim=0).to(torch.int32)
zero = torch.zeros(1, dtype=torch.int32, device="cuda")
return torch.cat((zero, cu_seqlens))


def _rotate_half(x: torch.Tensor) -> torch.Tensor:
"""
change sign so the last dimension becomes [-odd, +even]
Expand DownExpand Up@@ -345,7 +359,7 @@ def __init__(
_flash_attn_version >= _flash_attn_version_required
), f"FlashAttention minimum version {_flash_attn_version_required} is required."

self.attn_causal_mask = attn_mask_type == "causal"
self.attn_mask_type = attn_mask_type
self.norm_factor = norm_factor
self.attention_dropout_ctx = attention_dropout_ctx
self.attention_dropout = attention_dropout
Expand All@@ -356,6 +370,7 @@ def forward(
query_layer: torch.Tensor,
key_layer: torch.Tensor,
value_layer: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""flash-attn fprop"""

Expand All@@ -369,7 +384,6 @@ def forward(
), 'FlashAttention currently only supports CUDA tensors.'

# For now just 128, will make it more general in the future

if (query_layer.shape[-1] == 128 and
query_layer.shape[0] * query_layer.shape[1] >= 512 and
_check_if_interleaved_qkv(query_layer, key_layer, value_layer)):
Expand All@@ -389,18 +403,30 @@ def forward(
]

max_seqlen = seqlen
cu_seqlens = torch.arange(
0,
(batch_size + 1) * seqlen,
step=seqlen,
dtype=torch.int32,
device=query_layer.device)
if self.attn_mask_type == "padding":
assert (
attention_mask is not None
), "Boolean attention mask must be provided for padding."

@cyanguwacyanguwaJun 22, 2023

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.

So we're settled on requiring a mask tensor than a cu_seqlen tensor from users?

@ksivamanksivamanJun 22, 2023•
edited
Loading

Copy link
Copy Markdown
MemberAuthor

Choose a reason for hiding this comment

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

Yeah I think so, just to be consistent


mask_shape = torch.Size([seqlen, batch_size])
assert (
attention_mask.shape == mask_shape
), f"Expected shape {mask_shape} for attenion mask but found {attention_mask.shape}."

cu_seqlens = get_cu_seqlens(attention_mask)
else:
cu_seqlens = torch.arange(
0,
(batch_size + 1) * seqlen,
step=seqlen,
dtype=torch.int32,
device=query_layer.device)

with self.attention_dropout_ctx():
output = flash_attn_unpadded_func(
query_layer, key_layer, value_layer, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen,
self.attention_dropout if self.training else 0.0,
softmax_scale=1.0/self.norm_factor, causal=self.attn_causal_mask,
softmax_scale=1.0/self.norm_factor, causal=self.attn_mask_type=="causal",
deterministic=self.deterministic,
)

Expand DownExpand Up@@ -696,7 +722,7 @@ class DotProductAttention(torch.nn.Module):
.. note::

Argument :attr:`attention_mask` will be ignored in the `forward` call when
:attr:`attn_mask_type` is set to `"causal"`.
:attr:`attn_mask_type` is set to `"causal"` or `"no_mask"`.

.. warning::

Expand All@@ -714,7 +740,7 @@ class DotProductAttention(torch.nn.Module):
number of key-value channels.
attention_dropout: float, default = 0.0
dropout probability for the dropout op during multi-head attention.
attn_mask_type: {'causal', 'padding'}, default = `causal`
attn_mask_type: {'causal', 'padding', 'no_mask'}, default = `causal`
type of attention mask passed into softmax operation.
layer_number: int, default = `None`
layer number of the current `DotProductAttention` when multiple such modules
Expand DownExpand Up@@ -829,7 +855,7 @@ def forward(
.. note::

Argument :attr:`attention_mask` will be ignored when :attr:`attn_mask_type`
is set to `"causal"`.
is set to `"causal"` or `"no_mask"`.

.. note::

Expand DownExpand Up@@ -884,7 +910,6 @@ def forward(
use_flash_attention = False

if self.attn_mask_type == "padding" and attention_mask is not None:
use_flash_attention = False
use_fused_attention = False

if is_in_onnx_export_mode():
Expand All@@ -911,8 +936,9 @@ def forward(
return self._checkpointed_attention_forward(self.flash_attention,
query_layer,
key_layer,
value_layer)
return self.flash_attention(query_layer, key_layer, value_layer)
value_layer,
attention_mask)
return self.flash_attention(query_layer, key_layer, value_layer, attention_mask)

if use_fused_attention:
if checkpoint_core_attention:
Expand DownExpand Up@@ -1139,11 +1165,6 @@ def forward(
"""MultiHeadAttention FWD"""
# hidden_states: [sq, b, h]

if self.attn_mask_type != "causal" and attention_mask is not None:
assert (
attention_mask.dtype == torch.bool
), "Attention mask must be a boolean tensor"

# =================================================
# Pre-allocate memory for key-values for inference.
# =================================================
Expand Down
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { // Add copy buttons to all 
 blocks
(function() {
function addCopyButtons() {
document.querySelectorAll('pre code').forEach(function(codeBlock) {
if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;
codeBlock.parentElement.setAttribute('data-copy-added', 'true');
var btn = document.createElement('button');
btn.textContent = 'Copy';
btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';
btn.onmouseover = function() { this.style.opacity = '1'; };
btn.onmouseout = function() { this.style.opacity = '0.7'; };
btn.onclick = function() {
navigator.clipboard.writeText(codeBlock.textContent).then(function() {
btn.textContent = 'Copied!';
setTimeout(function() { btn.textContent = 'Copy'; }, 1500);
});
};
codeBlock.parentElement.style.position = 'relative';
codeBlock.parentElement.appendChild(btn);
});
}
addCopyButtons();
// Re-run on dynamic content
var observer = new MutationObserver(addCopyButtons);
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '
Expand FA coverage to padding masks by ksivaman · Pull Request #291 · NVIDIA/TransformerEngine · GitHub
Skip to content
Expand FA coverage to padding masks#291

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

 Closed

ksivaman
wants to merge
12
commits into

from

 Closed


Expand FA coverage to padding masks
#291
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


12 commits


Select commit
Hold shift + click to select a range
9a28af3

initial changes [wip]


ksivaman Jun 21, 2023
7b01a27

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 21, 2023
a89e43c

Add padding mask support for FA


ksivaman Jun 21, 2023
ae104b3

Address review comments


ksivaman Jun 21, 2023
e38d18d

rm causal mask from tests and add padding


ksivaman Jun 22, 2023
b0212bd

fix some conflicts


ksivaman Jun 22, 2023
e3dfd8b

conflicts


ksivaman Jun 22, 2023
ad8fc31

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 22, 2023
e2b2ddf

Fixes


ksivaman Jun 22, 2023
8eb8d90

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 23, 2023
e0cd408

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 26, 2023
e181188

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 27, 2023

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


13 changes: 5 additions & 8 deletions
tests/pytorch/test_numerics.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -401,11 +401,10 @@ def _test_e2e_selective_recompute(block, bs, dtype, config, recompute=False):
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_hidden_states.retain_grad()
 te_inp_attn_mask = get_causal_attn_mask(config.seq_len)

 te_out = block(
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 checkpoint_core_attention=recompute,
 )
 loss = te_out.sum()
Expand DownExpand Up@@ -461,7 +460,6 @@ def _test_e2e_full_recompute(block, bs, dtype, config, recompute=False):
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_hidden_states.retain_grad()
 te_inp_attn_mask = get_causal_attn_mask(config.seq_len)

 if recompute:
 te_out = te_checkpoint(
Expand All@@ -470,13 +468,13 @@ def _test_e2e_full_recompute(block, bs, dtype, config, recompute=False):
 get_dummy_cuda_rng_tracker,
 None, # tp_group
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 checkpoint_core_attention=False,
 )
 else:
 te_out = block(
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 checkpoint_core_attention=False,
 )
 loss = te_out.sum()
Expand DownExpand Up@@ -556,14 +554,13 @@ def _test_e2e_checkpointing(bs, dtype, config, checkpoint=False, steps=10, path=
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_hidden_states.retain_grad()
 te_inp_attn_mask = get_causal_attn_mask(config.seq_len)

 block = _test_e2e_checkpointing_get_model(config, dtype)

 for _ in range(steps // 2):
 te_out = block(
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 )
 loss = te_out.sum()
 loss.backward()
Expand DownExpand Up@@ -594,7 +591,7 @@ def _test_e2e_checkpointing(bs, dtype, config, checkpoint=False, steps=10, path=
 for _ in range(steps // 2):
 te_out = block(
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 )
 loss = te_out.sum()
 loss.backward()
Expand Down

91 changes: 39 additions & 52 deletions
tests/pytorch/test_sanity.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -154,26 +154,14 @@ def _test_sanity_e2e_amp(block, bs, dtype, config, fp8_recipe, skip_wgrad):
 config.seq_len, bs, config.hidden_size, dtype=torch.float32, requires_grad=True
 ).cuda()
 te_inp_hidden_states.retain_grad()
 te_inp_attn_mask = (
 torch.rand(
 (
 1,
 1,
 config.seq_len,
 config.seq_len,
 )
 )
 .cuda()
 .bool()
 )

 if skip_wgrad:
 _disable_wgrads(block)

 use_fp8 = fp8_recipe is not None
 with torch.autocast(device_type="cuda", enabled=True, dtype=dtype):
 with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
 te_out = block(te_inp_hidden_states, te_inp_attn_mask)
 te_out = block(te_inp_hidden_states)
 loss = te_out.sum()

 loss.backward()
Expand All@@ -190,18 +178,6 @@ def _test_sanity_e2e_gradient_accumulation_fusion(block, bs, dtype, config, fp8_
 te_inp_hidden_states = torch.randn(
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_attn_mask = (
 torch.rand(
 (
 1,
 1,
 config.seq_len,
 config.seq_len,
 )
 )
 .cuda()
 .bool()
 )

 if skip_wgrad:
 _disable_wgrads(block)
Expand All@@ -214,7 +190,7 @@ def _test_sanity_e2e_gradient_accumulation_fusion(block, bs, dtype, config, fp8_

 use_fp8 = fp8_recipe is not None
 with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
 te_out = block(te_inp_hidden_states, te_inp_attn_mask)
 te_out = block(te_inp_hidden_states)
 loss = te_out.sum()
 loss.backward()
 torch.cuda.synchronize()
Expand All@@ -232,18 +208,29 @@ def _test_sanity_e2e(block, bs, dtype, config, fp8_recipe, skip_wgrad):
 te_inp_hidden_states = torch.randn(
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_attn_mask = (
 torch.rand(
 (
 1,
 1,
 config.seq_len,
 config.seq_len,
 )
 )
 .cuda()
 .bool()
 )

 if skip_wgrad:
 _disable_wgrads(block)

 use_fp8 = fp8_recipe is not None
 with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
 te_out = block(te_inp_hidden_states)
 loss = te_out.sum()
 loss.backward()
 torch.cuda.synchronize()


def _test_sanity_e2e_bert(block, bs, dtype, config, fp8_recipe, skip_wgrad):
 te_inp_hidden_states = torch.randn(
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()

 if dtype == torch.float32:
 mask_shape = torch.Size([1, 1, config.seq_len, config.seq_len])
 else:
 mask_shape = torch.Size([config.seq_len, bs])

 te_inp_attn_mask = torch.rand(mask_shape).cuda().bool()

Comment on lines

+228
 to 
+233

@timmoon10timmoon10Jun 21, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Why does the FP32 test require a 4D mask while other dtypes use a 2D mask?
@ksivamanksivamanJun 21, 2023

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

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

The FP32 doesn't take the FA path and so uses the PyTorch torch softmax path for padding mask which expects this
@timmoon10timmoon10Jun 22, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Ah, so FA vs PyT is no longer an implementation detail since it expects a different mask format. This makes me think we need an option to explicitly enable or disable FA, and to error out instead of falling back to PyT.
@timmoon10timmoon10Jun 22, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Alternatively, we should expect the mask to be in the format for PyT, and then convert it to the FA format internally so it isn't visible to the user.
@ptrendxptrendxJun 23, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Agree with Tim. OR, accept both in both cases and generate the right one for the current implementation.
To be honest though, the current API for mask type is just bad. For example, what if you want padding AND causal? Then your only option is set the padding mask type and create the causal mask yourself (which would mean you can't do FA in the current implementation which is bad). There is also nothing stopping you really from having arbitrary padding mask (as in, with random elements zeroed out), which would break the assumptions you have in this PR. I think what we could do is introduce new names for the padding type (let's say "pad", "arbitrary" and "no_mask"), deprecate the old names and add causal as separate switch. Then we could say that pad only accepts the "mask" that is actually just list of sequence lengths, arbitrary is whatever you want (and will not go through FA) and causal could be switched irrespective of the mask type.

 if skip_wgrad:
 _disable_wgrads(block)
Expand All@@ -260,26 +247,24 @@ def _test_sanity_e2e_T5(block, bs, dtype, config, fp8_recipe, skip_wgrad):
 te_inp_hidden_states = torch.randn(
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_attn_mask = (
 torch.rand(
 (
 1,
 1,
 config.seq_len,
 config.seq_len,
 )
 )
 .cuda()
 .bool()
 )

 if dtype == torch.float32:
 mask_shape = torch.Size([1, 1, config.seq_len, config.seq_len])
 else:
 mask_shape = torch.Size([config.seq_len, bs])

 enc_dec_attn_mask = torch.rand(mask_shape).cuda().bool()

 if skip_wgrad:
 _disable_wgrads(block)

 use_fp8 = fp8_recipe is not None
 with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
 te_out = block(
 te_inp_hidden_states, te_inp_attn_mask, encoder_output=te_inp_hidden_states
 te_inp_hidden_states,
 None,
 encoder_output=te_inp_hidden_states,
 enc_dec_attn_mask=enc_dec_attn_mask,
 )
 loss = te_out.sum()
 loss.backward()
Expand DownExpand Up@@ -468,12 +453,14 @@ def test_sanity_bert(dtype, bs, fp8_recipe, model, skip_wgrad, zero_centered_gam
 apply_residual_connection_post_layernorm=True,
 output_layernorm=True,
 zero_centered_gamma=zero_centered_gamma,
 self_attn_mask_type="padding",

 )
 .to(dtype=dtype)
 .cuda()
 )

_test_sanity_e2e(block, bs, dtype, config, fp8_recipe, skip_wgrad)
_test_sanity_e2e_bert(block, bs, dtype, config, fp8_recipe, skip_wgrad)


@pytest.mark.parametrize("dtype", param_types)
Expand Down

61 changes: 41 additions & 20 deletions
transformer_engine/pytorch/attention.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -44,6 +44,7 @@
 checkpoint,
)
from transformer_engine.pytorch.export import is_in_onnx_export_mode
from transformer_engine.pytorch.jit import jit_fuser

_flash_attn_version = packaging.version.Version(version("flash-attn"))
_flash_attn_version_required = packaging.version.Version("1.0.6")
Expand All@@ -52,6 +53,19 @@
__all__ = ["DotProductAttention"]


@jit_fuser
def get_cu_seqlens(padding_mask: torch.Tensor) -> torch.Tensor:
 """
 Given a padding mask of shape [seq_len, batch_size], returns an int32
@cyanguwacyanguwaJun 22, 2023

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 guess by seq_len you mean max_seqlen? Also, should the mask tensor be required to be a CUDA tensor?
@ksivamanksivamanJun 22, 2023

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

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

Yeah it is the max_seqlen.
@ksivamanksivamanJun 22, 2023

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

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

Yeah it should be cuda tensor since the inp is cuda as well.
 tensor of shape [batch_size + 1,] containing the cumulative sequence
 lengths of every sample in the batch.
 """
 reduced_mask = padding_mask.sum(dim=0)
 cu_seqlens = reduced_mask.cumsum(dim=0).to(torch.int32)
 zero = torch.zeros(1, dtype=torch.int32, device="cuda")
 return torch.cat((zero, cu_seqlens))


def _rotate_half(x: torch.Tensor) -> torch.Tensor:
 """
 change sign so the last dimension becomes [-odd, +even]
Expand DownExpand Up@@ -345,7 +359,7 @@ def __init__(
 _flash_attn_version >= _flash_attn_version_required
 ), f"FlashAttention minimum version {_flash_attn_version_required} is required."

 self.attn_causal_mask = attn_mask_type == "causal"
 self.attn_mask_type = attn_mask_type
 self.norm_factor = norm_factor
 self.attention_dropout_ctx = attention_dropout_ctx
 self.attention_dropout = attention_dropout
Expand All@@ -356,6 +370,7 @@ def forward(
 query_layer: torch.Tensor,
 key_layer: torch.Tensor,
 value_layer: torch.Tensor,
 attention_mask: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
 """flash-attn fprop"""

Expand All@@ -369,7 +384,6 @@ def forward(
 ), 'FlashAttention currently only supports CUDA tensors.'

 # For now just 128, will make it more general in the future

 if (query_layer.shape[-1] == 128 and
 query_layer.shape[0] * query_layer.shape[1] >= 512 and
 _check_if_interleaved_qkv(query_layer, key_layer, value_layer)):
Expand All@@ -389,18 +403,30 @@ def forward(
 ]

 max_seqlen = seqlen
 cu_seqlens = torch.arange(
 0,
 (batch_size + 1) * seqlen,
 step=seqlen,
 dtype=torch.int32,
 device=query_layer.device)
 if self.attn_mask_type == "padding":
 assert (
 attention_mask is not None
 ), "Boolean attention mask must be provided for padding."
@cyanguwacyanguwaJun 22, 2023

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.

So we're settled on requiring a mask tensor than a cu_seqlen tensor from users?
@ksivamanksivamanJun 22, 2023•

edited

Loading

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

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

Yeah I think so, just to be consistent

 mask_shape = torch.Size([seqlen, batch_size])
 assert (
 attention_mask.shape == mask_shape
 ), f"Expected shape {mask_shape} for attenion mask but found {attention_mask.shape}."

 cu_seqlens = get_cu_seqlens(attention_mask)
 else:
 cu_seqlens = torch.arange(
 0,
 (batch_size + 1) * seqlen,
 step=seqlen,
 dtype=torch.int32,
 device=query_layer.device)

 with self.attention_dropout_ctx():
 output = flash_attn_unpadded_func(
 query_layer, key_layer, value_layer, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen,
 self.attention_dropout if self.training else 0.0,
 softmax_scale=1.0/self.norm_factor, causal=self.attn_causal_mask,
 softmax_scale=1.0/self.norm_factor, causal=self.attn_mask_type=="causal",
 deterministic=self.deterministic,
 )

Expand DownExpand Up@@ -696,7 +722,7 @@ class DotProductAttention(torch.nn.Module):
 .. note::

 Argument :attr:`attention_mask` will be ignored in the `forward` call when
 :attr:`attn_mask_type` is set to `"causal"`.
 :attr:`attn_mask_type` is set to `"causal"` or `"no_mask"`.

 .. warning::

Expand All@@ -714,7 +740,7 @@ class DotProductAttention(torch.nn.Module):
 number of key-value channels.
 attention_dropout: float, default = 0.0
 dropout probability for the dropout op during multi-head attention.
 attn_mask_type: {'causal', 'padding'}, default = `causal`
 attn_mask_type: {'causal', 'padding', 'no_mask'}, default = `causal`
 type of attention mask passed into softmax operation.
 layer_number: int, default = `None`
 layer number of the current `DotProductAttention` when multiple such modules
Expand DownExpand Up@@ -829,7 +855,7 @@ def forward(
 .. note::

 Argument :attr:`attention_mask` will be ignored when :attr:`attn_mask_type`
 is set to `"causal"`.
 is set to `"causal"` or `"no_mask"`.

 .. note::

Expand DownExpand Up@@ -884,7 +910,6 @@ def forward(
 use_flash_attention = False

 if self.attn_mask_type == "padding" and attention_mask is not None:
 use_flash_attention = False
 use_fused_attention = False

 if is_in_onnx_export_mode():
Expand All@@ -911,8 +936,9 @@ def forward(
 return self._checkpointed_attention_forward(self.flash_attention,
 query_layer,
 key_layer,
 value_layer)
 return self.flash_attention(query_layer, key_layer, value_layer)
 value_layer,
 attention_mask)
 return self.flash_attention(query_layer, key_layer, value_layer, attention_mask)

 if use_fused_attention:
 if checkpoint_core_attention:
Expand DownExpand Up@@ -1139,11 +1165,6 @@ def forward(
 """MultiHeadAttention FWD"""
 # hidden_states: [sq, b, h]

 if self.attn_mask_type != "causal" and attention_mask is not None:
 assert (
 attention_mask.dtype == torch.bool
 ), "Attention mask must be a boolean tensor"

 # =================================================
 # Pre-allocate memory for key-values for inference.
 # =================================================
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Force GitHub README to respect dark mode
(function() {
var style = document.createElement('style');
style.textContent = '
.markdown-body {
color-scheme: dark light;
}
.markdown-body pre { background: #161b22 !important; }
.markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }
.markdown-body table th, .markdown-body table td { border-color: #30363d !important; }
.markdown-body img { background: #0d1117; }
.markdown-body blockquote { border-left-color: #8b949e; }
.markdown-body hr { border-color: #30363d; }
';
document.head.appendChild(style);
})();
}
} catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
Expand FA coverage to padding masks by ksivaman · Pull Request #291 · NVIDIA/TransformerEngine · GitHub
Skip to content
Expand FA coverage to padding masks#291

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

 Closed

ksivaman
wants to merge
12
commits into

from

 Closed


Expand FA coverage to padding masks
#291
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


12 commits


Select commit
Hold shift + click to select a range
9a28af3

initial changes [wip]


ksivaman Jun 21, 2023
7b01a27

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 21, 2023
a89e43c

Add padding mask support for FA


ksivaman Jun 21, 2023
ae104b3

Address review comments


ksivaman Jun 21, 2023
e38d18d

rm causal mask from tests and add padding


ksivaman Jun 22, 2023
b0212bd

fix some conflicts


ksivaman Jun 22, 2023
e3dfd8b

conflicts


ksivaman Jun 22, 2023
ad8fc31

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 22, 2023
e2b2ddf

Fixes


ksivaman Jun 22, 2023
8eb8d90

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 23, 2023
e0cd408

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 26, 2023
e181188

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 27, 2023

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


13 changes: 5 additions & 8 deletions
tests/pytorch/test_numerics.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -401,11 +401,10 @@ def _test_e2e_selective_recompute(block, bs, dtype, config, recompute=False):
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_hidden_states.retain_grad()
 te_inp_attn_mask = get_causal_attn_mask(config.seq_len)

 te_out = block(
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 checkpoint_core_attention=recompute,
 )
 loss = te_out.sum()
Expand DownExpand Up@@ -461,7 +460,6 @@ def _test_e2e_full_recompute(block, bs, dtype, config, recompute=False):
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_hidden_states.retain_grad()
 te_inp_attn_mask = get_causal_attn_mask(config.seq_len)

 if recompute:
 te_out = te_checkpoint(
Expand All@@ -470,13 +468,13 @@ def _test_e2e_full_recompute(block, bs, dtype, config, recompute=False):
 get_dummy_cuda_rng_tracker,
 None, # tp_group
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 checkpoint_core_attention=False,
 )
 else:
 te_out = block(
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 checkpoint_core_attention=False,
 )
 loss = te_out.sum()
Expand DownExpand Up@@ -556,14 +554,13 @@ def _test_e2e_checkpointing(bs, dtype, config, checkpoint=False, steps=10, path=
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_hidden_states.retain_grad()
 te_inp_attn_mask = get_causal_attn_mask(config.seq_len)

 block = _test_e2e_checkpointing_get_model(config, dtype)

 for _ in range(steps // 2):
 te_out = block(
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 )
 loss = te_out.sum()
 loss.backward()
Expand DownExpand Up@@ -594,7 +591,7 @@ def _test_e2e_checkpointing(bs, dtype, config, checkpoint=False, steps=10, path=
 for _ in range(steps // 2):
 te_out = block(
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 )
 loss = te_out.sum()
 loss.backward()
Expand Down

91 changes: 39 additions & 52 deletions
tests/pytorch/test_sanity.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -154,26 +154,14 @@ def _test_sanity_e2e_amp(block, bs, dtype, config, fp8_recipe, skip_wgrad):
 config.seq_len, bs, config.hidden_size, dtype=torch.float32, requires_grad=True
 ).cuda()
 te_inp_hidden_states.retain_grad()
 te_inp_attn_mask = (
 torch.rand(
 (
 1,
 1,
 config.seq_len,
 config.seq_len,
 )
 )
 .cuda()
 .bool()
 )

 if skip_wgrad:
 _disable_wgrads(block)

 use_fp8 = fp8_recipe is not None
 with torch.autocast(device_type="cuda", enabled=True, dtype=dtype):
 with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
 te_out = block(te_inp_hidden_states, te_inp_attn_mask)
 te_out = block(te_inp_hidden_states)
 loss = te_out.sum()

 loss.backward()
Expand All@@ -190,18 +178,6 @@ def _test_sanity_e2e_gradient_accumulation_fusion(block, bs, dtype, config, fp8_
 te_inp_hidden_states = torch.randn(
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_attn_mask = (
 torch.rand(
 (
 1,
 1,
 config.seq_len,
 config.seq_len,
 )
 )
 .cuda()
 .bool()
 )

 if skip_wgrad:
 _disable_wgrads(block)
Expand All@@ -214,7 +190,7 @@ def _test_sanity_e2e_gradient_accumulation_fusion(block, bs, dtype, config, fp8_

 use_fp8 = fp8_recipe is not None
 with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
 te_out = block(te_inp_hidden_states, te_inp_attn_mask)
 te_out = block(te_inp_hidden_states)
 loss = te_out.sum()
 loss.backward()
 torch.cuda.synchronize()
Expand All@@ -232,18 +208,29 @@ def _test_sanity_e2e(block, bs, dtype, config, fp8_recipe, skip_wgrad):
 te_inp_hidden_states = torch.randn(
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_attn_mask = (
 torch.rand(
 (
 1,
 1,
 config.seq_len,
 config.seq_len,
 )
 )
 .cuda()
 .bool()
 )

 if skip_wgrad:
 _disable_wgrads(block)

 use_fp8 = fp8_recipe is not None
 with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
 te_out = block(te_inp_hidden_states)
 loss = te_out.sum()
 loss.backward()
 torch.cuda.synchronize()


def _test_sanity_e2e_bert(block, bs, dtype, config, fp8_recipe, skip_wgrad):
 te_inp_hidden_states = torch.randn(
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()

 if dtype == torch.float32:
 mask_shape = torch.Size([1, 1, config.seq_len, config.seq_len])
 else:
 mask_shape = torch.Size([config.seq_len, bs])

 te_inp_attn_mask = torch.rand(mask_shape).cuda().bool()

Comment on lines

+228
 to 
+233

@timmoon10timmoon10Jun 21, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Why does the FP32 test require a 4D mask while other dtypes use a 2D mask?
@ksivamanksivamanJun 21, 2023

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

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

The FP32 doesn't take the FA path and so uses the PyTorch torch softmax path for padding mask which expects this
@timmoon10timmoon10Jun 22, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Ah, so FA vs PyT is no longer an implementation detail since it expects a different mask format. This makes me think we need an option to explicitly enable or disable FA, and to error out instead of falling back to PyT.
@timmoon10timmoon10Jun 22, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Alternatively, we should expect the mask to be in the format for PyT, and then convert it to the FA format internally so it isn't visible to the user.
@ptrendxptrendxJun 23, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Agree with Tim. OR, accept both in both cases and generate the right one for the current implementation.
To be honest though, the current API for mask type is just bad. For example, what if you want padding AND causal? Then your only option is set the padding mask type and create the causal mask yourself (which would mean you can't do FA in the current implementation which is bad). There is also nothing stopping you really from having arbitrary padding mask (as in, with random elements zeroed out), which would break the assumptions you have in this PR. I think what we could do is introduce new names for the padding type (let's say "pad", "arbitrary" and "no_mask"), deprecate the old names and add causal as separate switch. Then we could say that pad only accepts the "mask" that is actually just list of sequence lengths, arbitrary is whatever you want (and will not go through FA) and causal could be switched irrespective of the mask type.

 if skip_wgrad:
 _disable_wgrads(block)
Expand All@@ -260,26 +247,24 @@ def _test_sanity_e2e_T5(block, bs, dtype, config, fp8_recipe, skip_wgrad):
 te_inp_hidden_states = torch.randn(
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_attn_mask = (
 torch.rand(
 (
 1,
 1,
 config.seq_len,
 config.seq_len,
 )
 )
 .cuda()
 .bool()
 )

 if dtype == torch.float32:
 mask_shape = torch.Size([1, 1, config.seq_len, config.seq_len])
 else:
 mask_shape = torch.Size([config.seq_len, bs])

 enc_dec_attn_mask = torch.rand(mask_shape).cuda().bool()

 if skip_wgrad:
 _disable_wgrads(block)

 use_fp8 = fp8_recipe is not None
 with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
 te_out = block(
 te_inp_hidden_states, te_inp_attn_mask, encoder_output=te_inp_hidden_states
 te_inp_hidden_states,
 None,
 encoder_output=te_inp_hidden_states,
 enc_dec_attn_mask=enc_dec_attn_mask,
 )
 loss = te_out.sum()
 loss.backward()
Expand DownExpand Up@@ -468,12 +453,14 @@ def test_sanity_bert(dtype, bs, fp8_recipe, model, skip_wgrad, zero_centered_gam
 apply_residual_connection_post_layernorm=True,
 output_layernorm=True,
 zero_centered_gamma=zero_centered_gamma,
 self_attn_mask_type="padding",

 )
 .to(dtype=dtype)
 .cuda()
 )

_test_sanity_e2e(block, bs, dtype, config, fp8_recipe, skip_wgrad)
_test_sanity_e2e_bert(block, bs, dtype, config, fp8_recipe, skip_wgrad)


@pytest.mark.parametrize("dtype", param_types)
Expand Down

61 changes: 41 additions & 20 deletions
transformer_engine/pytorch/attention.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -44,6 +44,7 @@
 checkpoint,
)
from transformer_engine.pytorch.export import is_in_onnx_export_mode
from transformer_engine.pytorch.jit import jit_fuser

_flash_attn_version = packaging.version.Version(version("flash-attn"))
_flash_attn_version_required = packaging.version.Version("1.0.6")
Expand All@@ -52,6 +53,19 @@
__all__ = ["DotProductAttention"]


@jit_fuser
def get_cu_seqlens(padding_mask: torch.Tensor) -> torch.Tensor:
 """
 Given a padding mask of shape [seq_len, batch_size], returns an int32
@cyanguwacyanguwaJun 22, 2023

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 guess by seq_len you mean max_seqlen? Also, should the mask tensor be required to be a CUDA tensor?
@ksivamanksivamanJun 22, 2023

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

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

Yeah it is the max_seqlen.
@ksivamanksivamanJun 22, 2023

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

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

Yeah it should be cuda tensor since the inp is cuda as well.
 tensor of shape [batch_size + 1,] containing the cumulative sequence
 lengths of every sample in the batch.
 """
 reduced_mask = padding_mask.sum(dim=0)
 cu_seqlens = reduced_mask.cumsum(dim=0).to(torch.int32)
 zero = torch.zeros(1, dtype=torch.int32, device="cuda")
 return torch.cat((zero, cu_seqlens))


def _rotate_half(x: torch.Tensor) -> torch.Tensor:
 """
 change sign so the last dimension becomes [-odd, +even]
Expand DownExpand Up@@ -345,7 +359,7 @@ def __init__(
 _flash_attn_version >= _flash_attn_version_required
 ), f"FlashAttention minimum version {_flash_attn_version_required} is required."

 self.attn_causal_mask = attn_mask_type == "causal"
 self.attn_mask_type = attn_mask_type
 self.norm_factor = norm_factor
 self.attention_dropout_ctx = attention_dropout_ctx
 self.attention_dropout = attention_dropout
Expand All@@ -356,6 +370,7 @@ def forward(
 query_layer: torch.Tensor,
 key_layer: torch.Tensor,
 value_layer: torch.Tensor,
 attention_mask: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
 """flash-attn fprop"""

Expand All@@ -369,7 +384,6 @@ def forward(
 ), 'FlashAttention currently only supports CUDA tensors.'

 # For now just 128, will make it more general in the future

 if (query_layer.shape[-1] == 128 and
 query_layer.shape[0] * query_layer.shape[1] >= 512 and
 _check_if_interleaved_qkv(query_layer, key_layer, value_layer)):
Expand All@@ -389,18 +403,30 @@ def forward(
 ]

 max_seqlen = seqlen
 cu_seqlens = torch.arange(
 0,
 (batch_size + 1) * seqlen,
 step=seqlen,
 dtype=torch.int32,
 device=query_layer.device)
 if self.attn_mask_type == "padding":
 assert (
 attention_mask is not None
 ), "Boolean attention mask must be provided for padding."
@cyanguwacyanguwaJun 22, 2023

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.

So we're settled on requiring a mask tensor than a cu_seqlen tensor from users?
@ksivamanksivamanJun 22, 2023•

edited

Loading

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

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

Yeah I think so, just to be consistent

 mask_shape = torch.Size([seqlen, batch_size])
 assert (
 attention_mask.shape == mask_shape
 ), f"Expected shape {mask_shape} for attenion mask but found {attention_mask.shape}."

 cu_seqlens = get_cu_seqlens(attention_mask)
 else:
 cu_seqlens = torch.arange(
 0,
 (batch_size + 1) * seqlen,
 step=seqlen,
 dtype=torch.int32,
 device=query_layer.device)

 with self.attention_dropout_ctx():
 output = flash_attn_unpadded_func(
 query_layer, key_layer, value_layer, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen,
 self.attention_dropout if self.training else 0.0,
 softmax_scale=1.0/self.norm_factor, causal=self.attn_causal_mask,
 softmax_scale=1.0/self.norm_factor, causal=self.attn_mask_type=="causal",
 deterministic=self.deterministic,
 )

Expand DownExpand Up@@ -696,7 +722,7 @@ class DotProductAttention(torch.nn.Module):
 .. note::

 Argument :attr:`attention_mask` will be ignored in the `forward` call when
 :attr:`attn_mask_type` is set to `"causal"`.
 :attr:`attn_mask_type` is set to `"causal"` or `"no_mask"`.

 .. warning::

Expand All@@ -714,7 +740,7 @@ class DotProductAttention(torch.nn.Module):
 number of key-value channels.
 attention_dropout: float, default = 0.0
 dropout probability for the dropout op during multi-head attention.
 attn_mask_type: {'causal', 'padding'}, default = `causal`
 attn_mask_type: {'causal', 'padding', 'no_mask'}, default = `causal`
 type of attention mask passed into softmax operation.
 layer_number: int, default = `None`
 layer number of the current `DotProductAttention` when multiple such modules
Expand DownExpand Up@@ -829,7 +855,7 @@ def forward(
 .. note::

 Argument :attr:`attention_mask` will be ignored when :attr:`attn_mask_type`
 is set to `"causal"`.
 is set to `"causal"` or `"no_mask"`.

 .. note::

Expand DownExpand Up@@ -884,7 +910,6 @@ def forward(
 use_flash_attention = False

 if self.attn_mask_type == "padding" and attention_mask is not None:
 use_flash_attention = False
 use_fused_attention = False

 if is_in_onnx_export_mode():
Expand All@@ -911,8 +936,9 @@ def forward(
 return self._checkpointed_attention_forward(self.flash_attention,
 query_layer,
 key_layer,
 value_layer)
 return self.flash_attention(query_layer, key_layer, value_layer)
 value_layer,
 attention_mask)
 return self.flash_attention(query_layer, key_layer, value_layer, attention_mask)

 if use_fused_attention:
 if checkpoint_core_attention:
Expand DownExpand Up@@ -1139,11 +1165,6 @@ def forward(
 """MultiHeadAttention FWD"""
 # hidden_states: [sq, b, h]

 if self.attn_mask_type != "causal" and attention_mask is not None:
 assert (
 attention_mask.dtype == torch.bool
 ), "Attention mask must be a boolean tensor"

 # =================================================
 # Pre-allocate memory for key-values for inference.
 # =================================================
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Highlight search terms from Google/DuckDuckGo/Bing referrer
(function() {
var ref = document.referrer;
var terms = [];
if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {
var url = new URL(ref);
var q = url.searchParams.get('q') || url.searchParams.get('p');
if (q) {
terms = q.split(/\s+/).filter(function(t) { return t.length > 2; });
}
}
if (terms.length === 0) return;
var style = document.createElement('style');
style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';
document.head.appendChild(style);
function highlight(node) {
if (node.nodeType === 3) { // text node
var text = node.textContent;
var found = false;
terms.forEach(function(term) {
var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\]\\]/g, '\\') + ')', 'gi');
if (regex.test(text)) {
found = true;
var frag = document.createDocumentFragment();
var parts = text.split(regex);
parts.forEach(function(part, i) {
if (i % 2 === 0) {
frag.appendChild(document.createTextNode(part));
} else {
var span = document.createElement('span');
span.className = 'userscript-highlight';
span.textContent = part;
frag.appendChild(span);
}
});
node.parentNode.replaceChild(frag, node);
}
});
} else if (node.nodeType === 1 && node.childNodes) { // element
var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];
if (!skipTags.includes(node.tagName)) {
Array.from(node.childNodes).forEach(highlight);
}
}
}
highlight(document.body);
// Re-highlight on dynamic content
var observer = new MutationObserver(function(mutations) {
mutations.forEach(function(m) {
m.addedNodes.forEach(function(node) {
if (node.nodeType === 1 || node.nodeType === 3) highlight(node);
});
});
});
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
Expand FA coverage to padding masks by ksivaman · Pull Request #291 · NVIDIA/TransformerEngine · GitHub
Skip to content
Expand FA coverage to padding masks#291

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

 Closed

ksivaman
wants to merge
12
commits into

from

 Closed


Expand FA coverage to padding masks
#291
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


12 commits


Select commit
Hold shift + click to select a range
9a28af3

initial changes [wip]


ksivaman Jun 21, 2023
7b01a27

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 21, 2023
a89e43c

Add padding mask support for FA


ksivaman Jun 21, 2023
ae104b3

Address review comments


ksivaman Jun 21, 2023
e38d18d

rm causal mask from tests and add padding


ksivaman Jun 22, 2023
b0212bd

fix some conflicts


ksivaman Jun 22, 2023
e3dfd8b

conflicts


ksivaman Jun 22, 2023
ad8fc31

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 22, 2023
e2b2ddf

Fixes


ksivaman Jun 22, 2023
8eb8d90

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 23, 2023
e0cd408

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 26, 2023
e181188

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 27, 2023

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


13 changes: 5 additions & 8 deletions
tests/pytorch/test_numerics.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -401,11 +401,10 @@ def _test_e2e_selective_recompute(block, bs, dtype, config, recompute=False):
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_hidden_states.retain_grad()
 te_inp_attn_mask = get_causal_attn_mask(config.seq_len)

 te_out = block(
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 checkpoint_core_attention=recompute,
 )
 loss = te_out.sum()
Expand DownExpand Up@@ -461,7 +460,6 @@ def _test_e2e_full_recompute(block, bs, dtype, config, recompute=False):
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_hidden_states.retain_grad()
 te_inp_attn_mask = get_causal_attn_mask(config.seq_len)

 if recompute:
 te_out = te_checkpoint(
Expand All@@ -470,13 +468,13 @@ def _test_e2e_full_recompute(block, bs, dtype, config, recompute=False):
 get_dummy_cuda_rng_tracker,
 None, # tp_group
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 checkpoint_core_attention=False,
 )
 else:
 te_out = block(
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 checkpoint_core_attention=False,
 )
 loss = te_out.sum()
Expand DownExpand Up@@ -556,14 +554,13 @@ def _test_e2e_checkpointing(bs, dtype, config, checkpoint=False, steps=10, path=
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_hidden_states.retain_grad()
 te_inp_attn_mask = get_causal_attn_mask(config.seq_len)

 block = _test_e2e_checkpointing_get_model(config, dtype)

 for _ in range(steps // 2):
 te_out = block(
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 )
 loss = te_out.sum()
 loss.backward()
Expand DownExpand Up@@ -594,7 +591,7 @@ def _test_e2e_checkpointing(bs, dtype, config, checkpoint=False, steps=10, path=
 for _ in range(steps // 2):
 te_out = block(
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 )
 loss = te_out.sum()
 loss.backward()
Expand Down

91 changes: 39 additions & 52 deletions
tests/pytorch/test_sanity.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -154,26 +154,14 @@ def _test_sanity_e2e_amp(block, bs, dtype, config, fp8_recipe, skip_wgrad):
 config.seq_len, bs, config.hidden_size, dtype=torch.float32, requires_grad=True
 ).cuda()
 te_inp_hidden_states.retain_grad()
 te_inp_attn_mask = (
 torch.rand(
 (
 1,
 1,
 config.seq_len,
 config.seq_len,
 )
 )
 .cuda()
 .bool()
 )

 if skip_wgrad:
 _disable_wgrads(block)

 use_fp8 = fp8_recipe is not None
 with torch.autocast(device_type="cuda", enabled=True, dtype=dtype):
 with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
 te_out = block(te_inp_hidden_states, te_inp_attn_mask)
 te_out = block(te_inp_hidden_states)
 loss = te_out.sum()

 loss.backward()
Expand All@@ -190,18 +178,6 @@ def _test_sanity_e2e_gradient_accumulation_fusion(block, bs, dtype, config, fp8_
 te_inp_hidden_states = torch.randn(
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_attn_mask = (
 torch.rand(
 (
 1,
 1,
 config.seq_len,
 config.seq_len,
 )
 )
 .cuda()
 .bool()
 )

 if skip_wgrad:
 _disable_wgrads(block)
Expand All@@ -214,7 +190,7 @@ def _test_sanity_e2e_gradient_accumulation_fusion(block, bs, dtype, config, fp8_

 use_fp8 = fp8_recipe is not None
 with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
 te_out = block(te_inp_hidden_states, te_inp_attn_mask)
 te_out = block(te_inp_hidden_states)
 loss = te_out.sum()
 loss.backward()
 torch.cuda.synchronize()
Expand All@@ -232,18 +208,29 @@ def _test_sanity_e2e(block, bs, dtype, config, fp8_recipe, skip_wgrad):
 te_inp_hidden_states = torch.randn(
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_attn_mask = (
 torch.rand(
 (
 1,
 1,
 config.seq_len,
 config.seq_len,
 )
 )
 .cuda()
 .bool()
 )

 if skip_wgrad:
 _disable_wgrads(block)

 use_fp8 = fp8_recipe is not None
 with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
 te_out = block(te_inp_hidden_states)
 loss = te_out.sum()
 loss.backward()
 torch.cuda.synchronize()


def _test_sanity_e2e_bert(block, bs, dtype, config, fp8_recipe, skip_wgrad):
 te_inp_hidden_states = torch.randn(
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()

 if dtype == torch.float32:
 mask_shape = torch.Size([1, 1, config.seq_len, config.seq_len])
 else:
 mask_shape = torch.Size([config.seq_len, bs])

 te_inp_attn_mask = torch.rand(mask_shape).cuda().bool()

Comment on lines

+228
 to 
+233

@timmoon10timmoon10Jun 21, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Why does the FP32 test require a 4D mask while other dtypes use a 2D mask?
@ksivamanksivamanJun 21, 2023

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

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

The FP32 doesn't take the FA path and so uses the PyTorch torch softmax path for padding mask which expects this
@timmoon10timmoon10Jun 22, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Ah, so FA vs PyT is no longer an implementation detail since it expects a different mask format. This makes me think we need an option to explicitly enable or disable FA, and to error out instead of falling back to PyT.
@timmoon10timmoon10Jun 22, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Alternatively, we should expect the mask to be in the format for PyT, and then convert it to the FA format internally so it isn't visible to the user.
@ptrendxptrendxJun 23, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Agree with Tim. OR, accept both in both cases and generate the right one for the current implementation.
To be honest though, the current API for mask type is just bad. For example, what if you want padding AND causal? Then your only option is set the padding mask type and create the causal mask yourself (which would mean you can't do FA in the current implementation which is bad). There is also nothing stopping you really from having arbitrary padding mask (as in, with random elements zeroed out), which would break the assumptions you have in this PR. I think what we could do is introduce new names for the padding type (let's say "pad", "arbitrary" and "no_mask"), deprecate the old names and add causal as separate switch. Then we could say that pad only accepts the "mask" that is actually just list of sequence lengths, arbitrary is whatever you want (and will not go through FA) and causal could be switched irrespective of the mask type.

 if skip_wgrad:
 _disable_wgrads(block)
Expand All@@ -260,26 +247,24 @@ def _test_sanity_e2e_T5(block, bs, dtype, config, fp8_recipe, skip_wgrad):
 te_inp_hidden_states = torch.randn(
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_attn_mask = (
 torch.rand(
 (
 1,
 1,
 config.seq_len,
 config.seq_len,
 )
 )
 .cuda()
 .bool()
 )

 if dtype == torch.float32:
 mask_shape = torch.Size([1, 1, config.seq_len, config.seq_len])
 else:
 mask_shape = torch.Size([config.seq_len, bs])

 enc_dec_attn_mask = torch.rand(mask_shape).cuda().bool()

 if skip_wgrad:
 _disable_wgrads(block)

 use_fp8 = fp8_recipe is not None
 with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
 te_out = block(
 te_inp_hidden_states, te_inp_attn_mask, encoder_output=te_inp_hidden_states
 te_inp_hidden_states,
 None,
 encoder_output=te_inp_hidden_states,
 enc_dec_attn_mask=enc_dec_attn_mask,
 )
 loss = te_out.sum()
 loss.backward()
Expand DownExpand Up@@ -468,12 +453,14 @@ def test_sanity_bert(dtype, bs, fp8_recipe, model, skip_wgrad, zero_centered_gam
 apply_residual_connection_post_layernorm=True,
 output_layernorm=True,
 zero_centered_gamma=zero_centered_gamma,
 self_attn_mask_type="padding",

 )
 .to(dtype=dtype)
 .cuda()
 )

_test_sanity_e2e(block, bs, dtype, config, fp8_recipe, skip_wgrad)
_test_sanity_e2e_bert(block, bs, dtype, config, fp8_recipe, skip_wgrad)


@pytest.mark.parametrize("dtype", param_types)
Expand Down

61 changes: 41 additions & 20 deletions
transformer_engine/pytorch/attention.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -44,6 +44,7 @@
 checkpoint,
)
from transformer_engine.pytorch.export import is_in_onnx_export_mode
from transformer_engine.pytorch.jit import jit_fuser

_flash_attn_version = packaging.version.Version(version("flash-attn"))
_flash_attn_version_required = packaging.version.Version("1.0.6")
Expand All@@ -52,6 +53,19 @@
__all__ = ["DotProductAttention"]


@jit_fuser
def get_cu_seqlens(padding_mask: torch.Tensor) -> torch.Tensor:
 """
 Given a padding mask of shape [seq_len, batch_size], returns an int32
@cyanguwacyanguwaJun 22, 2023

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 guess by seq_len you mean max_seqlen? Also, should the mask tensor be required to be a CUDA tensor?
@ksivamanksivamanJun 22, 2023

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

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

Yeah it is the max_seqlen.
@ksivamanksivamanJun 22, 2023

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

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

Yeah it should be cuda tensor since the inp is cuda as well.
 tensor of shape [batch_size + 1,] containing the cumulative sequence
 lengths of every sample in the batch.
 """
 reduced_mask = padding_mask.sum(dim=0)
 cu_seqlens = reduced_mask.cumsum(dim=0).to(torch.int32)
 zero = torch.zeros(1, dtype=torch.int32, device="cuda")
 return torch.cat((zero, cu_seqlens))


def _rotate_half(x: torch.Tensor) -> torch.Tensor:
 """
 change sign so the last dimension becomes [-odd, +even]
Expand DownExpand Up@@ -345,7 +359,7 @@ def __init__(
 _flash_attn_version >= _flash_attn_version_required
 ), f"FlashAttention minimum version {_flash_attn_version_required} is required."

 self.attn_causal_mask = attn_mask_type == "causal"
 self.attn_mask_type = attn_mask_type
 self.norm_factor = norm_factor
 self.attention_dropout_ctx = attention_dropout_ctx
 self.attention_dropout = attention_dropout
Expand All@@ -356,6 +370,7 @@ def forward(
 query_layer: torch.Tensor,
 key_layer: torch.Tensor,
 value_layer: torch.Tensor,
 attention_mask: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
 """flash-attn fprop"""

Expand All@@ -369,7 +384,6 @@ def forward(
 ), 'FlashAttention currently only supports CUDA tensors.'

 # For now just 128, will make it more general in the future

 if (query_layer.shape[-1] == 128 and
 query_layer.shape[0] * query_layer.shape[1] >= 512 and
 _check_if_interleaved_qkv(query_layer, key_layer, value_layer)):
Expand All@@ -389,18 +403,30 @@ def forward(
 ]

 max_seqlen = seqlen
 cu_seqlens = torch.arange(
 0,
 (batch_size + 1) * seqlen,
 step=seqlen,
 dtype=torch.int32,
 device=query_layer.device)
 if self.attn_mask_type == "padding":
 assert (
 attention_mask is not None
 ), "Boolean attention mask must be provided for padding."
@cyanguwacyanguwaJun 22, 2023

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.

So we're settled on requiring a mask tensor than a cu_seqlen tensor from users?
@ksivamanksivamanJun 22, 2023•

edited

Loading

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

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

Yeah I think so, just to be consistent

 mask_shape = torch.Size([seqlen, batch_size])
 assert (
 attention_mask.shape == mask_shape
 ), f"Expected shape {mask_shape} for attenion mask but found {attention_mask.shape}."

 cu_seqlens = get_cu_seqlens(attention_mask)
 else:
 cu_seqlens = torch.arange(
 0,
 (batch_size + 1) * seqlen,
 step=seqlen,
 dtype=torch.int32,
 device=query_layer.device)

 with self.attention_dropout_ctx():
 output = flash_attn_unpadded_func(
 query_layer, key_layer, value_layer, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen,
 self.attention_dropout if self.training else 0.0,
 softmax_scale=1.0/self.norm_factor, causal=self.attn_causal_mask,
 softmax_scale=1.0/self.norm_factor, causal=self.attn_mask_type=="causal",
 deterministic=self.deterministic,
 )

Expand DownExpand Up@@ -696,7 +722,7 @@ class DotProductAttention(torch.nn.Module):
 .. note::

 Argument :attr:`attention_mask` will be ignored in the `forward` call when
 :attr:`attn_mask_type` is set to `"causal"`.
 :attr:`attn_mask_type` is set to `"causal"` or `"no_mask"`.

 .. warning::

Expand All@@ -714,7 +740,7 @@ class DotProductAttention(torch.nn.Module):
 number of key-value channels.
 attention_dropout: float, default = 0.0
 dropout probability for the dropout op during multi-head attention.
 attn_mask_type: {'causal', 'padding'}, default = `causal`
 attn_mask_type: {'causal', 'padding', 'no_mask'}, default = `causal`
 type of attention mask passed into softmax operation.
 layer_number: int, default = `None`
 layer number of the current `DotProductAttention` when multiple such modules
Expand DownExpand Up@@ -829,7 +855,7 @@ def forward(
 .. note::

 Argument :attr:`attention_mask` will be ignored when :attr:`attn_mask_type`
 is set to `"causal"`.
 is set to `"causal"` or `"no_mask"`.

 .. note::

Expand DownExpand Up@@ -884,7 +910,6 @@ def forward(
 use_flash_attention = False

 if self.attn_mask_type == "padding" and attention_mask is not None:
 use_flash_attention = False
 use_fused_attention = False

 if is_in_onnx_export_mode():
Expand All@@ -911,8 +936,9 @@ def forward(
 return self._checkpointed_attention_forward(self.flash_attention,
 query_layer,
 key_layer,
 value_layer)
 return self.flash_attention(query_layer, key_layer, value_layer)
 value_layer,
 attention_mask)
 return self.flash_attention(query_layer, key_layer, value_layer, attention_mask)

 if use_fused_attention:
 if checkpoint_core_attention:
Expand DownExpand Up@@ -1139,11 +1165,6 @@ def forward(
 """MultiHeadAttention FWD"""
 # hidden_states: [sq, b, h]

 if self.attn_mask_type != "causal" and attention_mask is not None:
 assert (
 attention_mask.dtype == torch.bool
 ), "Attention mask must be a boolean tensor"

 # =================================================
 # Pre-allocate memory for key-values for inference.
 # =================================================
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Strip utm_, fbclid, gclid, etc. from all links on page
(function() {
var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',
'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',
'ref', 'ref_src', 'source', 'medium', 'campaign'];
function cleanUrl(url) {
try {
var u = new URL(url, window.location.origin);
var changed = false;
trackingParams.forEach(function(p) {
if (u.searchParams.has(p)) {
u.searchParams.delete(p);
changed = true;
}
});
return changed ? u.toString() : url;
} catch (e) {
return url;
}
}
function cleanLinks() {
document.querySelectorAll('a[href]').forEach(function(a) {
var clean = cleanUrl(a.href);
if (clean !== a.href) a.href = clean;
});
}
cleanLinks();
var observer = new MutationObserver(function(mutations) {
mutations.forEach(function(m) {
m.addedNodes.forEach(function(node) {
if (node.nodeType === 1) {
if (node.tagName === 'A') cleanLinks();
node.querySelectorAll('a[href]').forEach(function(a) {
var clean = cleanUrl(a.href);
if (clean !== a.href) a.href = clean;
});
}
});
});
});
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); }
})();
(function(){
try {
var __m = "youtube.com";
var __re = new RegExp('^' + "youtube\\.com" + '
Expand FA coverage to padding masks by ksivaman · Pull Request #291 · NVIDIA/TransformerEngine · GitHub
Skip to content
Expand FA coverage to padding masks#291

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

 Closed

ksivaman
wants to merge
12
commits into

from

 Closed


Expand FA coverage to padding masks
#291
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


12 commits


Select commit
Hold shift + click to select a range
9a28af3

initial changes [wip]


ksivaman Jun 21, 2023
7b01a27

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 21, 2023
a89e43c

Add padding mask support for FA


ksivaman Jun 21, 2023
ae104b3

Address review comments


ksivaman Jun 21, 2023
e38d18d

rm causal mask from tests and add padding


ksivaman Jun 22, 2023
b0212bd

fix some conflicts


ksivaman Jun 22, 2023
e3dfd8b

conflicts


ksivaman Jun 22, 2023
ad8fc31

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 22, 2023
e2b2ddf

Fixes


ksivaman Jun 22, 2023
8eb8d90

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 23, 2023
e0cd408

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 26, 2023
e181188

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 27, 2023

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


13 changes: 5 additions & 8 deletions
tests/pytorch/test_numerics.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -401,11 +401,10 @@ def _test_e2e_selective_recompute(block, bs, dtype, config, recompute=False):
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_hidden_states.retain_grad()
 te_inp_attn_mask = get_causal_attn_mask(config.seq_len)

 te_out = block(
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 checkpoint_core_attention=recompute,
 )
 loss = te_out.sum()
Expand DownExpand Up@@ -461,7 +460,6 @@ def _test_e2e_full_recompute(block, bs, dtype, config, recompute=False):
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_hidden_states.retain_grad()
 te_inp_attn_mask = get_causal_attn_mask(config.seq_len)

 if recompute:
 te_out = te_checkpoint(
Expand All@@ -470,13 +468,13 @@ def _test_e2e_full_recompute(block, bs, dtype, config, recompute=False):
 get_dummy_cuda_rng_tracker,
 None, # tp_group
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 checkpoint_core_attention=False,
 )
 else:
 te_out = block(
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 checkpoint_core_attention=False,
 )
 loss = te_out.sum()
Expand DownExpand Up@@ -556,14 +554,13 @@ def _test_e2e_checkpointing(bs, dtype, config, checkpoint=False, steps=10, path=
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_hidden_states.retain_grad()
 te_inp_attn_mask = get_causal_attn_mask(config.seq_len)

 block = _test_e2e_checkpointing_get_model(config, dtype)

 for _ in range(steps // 2):
 te_out = block(
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 )
 loss = te_out.sum()
 loss.backward()
Expand DownExpand Up@@ -594,7 +591,7 @@ def _test_e2e_checkpointing(bs, dtype, config, checkpoint=False, steps=10, path=
 for _ in range(steps // 2):
 te_out = block(
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 )
 loss = te_out.sum()
 loss.backward()
Expand Down

91 changes: 39 additions & 52 deletions
tests/pytorch/test_sanity.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -154,26 +154,14 @@ def _test_sanity_e2e_amp(block, bs, dtype, config, fp8_recipe, skip_wgrad):
 config.seq_len, bs, config.hidden_size, dtype=torch.float32, requires_grad=True
 ).cuda()
 te_inp_hidden_states.retain_grad()
 te_inp_attn_mask = (
 torch.rand(
 (
 1,
 1,
 config.seq_len,
 config.seq_len,
 )
 )
 .cuda()
 .bool()
 )

 if skip_wgrad:
 _disable_wgrads(block)

 use_fp8 = fp8_recipe is not None
 with torch.autocast(device_type="cuda", enabled=True, dtype=dtype):
 with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
 te_out = block(te_inp_hidden_states, te_inp_attn_mask)
 te_out = block(te_inp_hidden_states)
 loss = te_out.sum()

 loss.backward()
Expand All@@ -190,18 +178,6 @@ def _test_sanity_e2e_gradient_accumulation_fusion(block, bs, dtype, config, fp8_
 te_inp_hidden_states = torch.randn(
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_attn_mask = (
 torch.rand(
 (
 1,
 1,
 config.seq_len,
 config.seq_len,
 )
 )
 .cuda()
 .bool()
 )

 if skip_wgrad:
 _disable_wgrads(block)
Expand All@@ -214,7 +190,7 @@ def _test_sanity_e2e_gradient_accumulation_fusion(block, bs, dtype, config, fp8_

 use_fp8 = fp8_recipe is not None
 with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
 te_out = block(te_inp_hidden_states, te_inp_attn_mask)
 te_out = block(te_inp_hidden_states)
 loss = te_out.sum()
 loss.backward()
 torch.cuda.synchronize()
Expand All@@ -232,18 +208,29 @@ def _test_sanity_e2e(block, bs, dtype, config, fp8_recipe, skip_wgrad):
 te_inp_hidden_states = torch.randn(
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_attn_mask = (
 torch.rand(
 (
 1,
 1,
 config.seq_len,
 config.seq_len,
 )
 )
 .cuda()
 .bool()
 )

 if skip_wgrad:
 _disable_wgrads(block)

 use_fp8 = fp8_recipe is not None
 with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
 te_out = block(te_inp_hidden_states)
 loss = te_out.sum()
 loss.backward()
 torch.cuda.synchronize()


def _test_sanity_e2e_bert(block, bs, dtype, config, fp8_recipe, skip_wgrad):
 te_inp_hidden_states = torch.randn(
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()

 if dtype == torch.float32:
 mask_shape = torch.Size([1, 1, config.seq_len, config.seq_len])
 else:
 mask_shape = torch.Size([config.seq_len, bs])

 te_inp_attn_mask = torch.rand(mask_shape).cuda().bool()

Comment on lines

+228
 to 
+233

@timmoon10timmoon10Jun 21, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Why does the FP32 test require a 4D mask while other dtypes use a 2D mask?
@ksivamanksivamanJun 21, 2023

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

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

The FP32 doesn't take the FA path and so uses the PyTorch torch softmax path for padding mask which expects this
@timmoon10timmoon10Jun 22, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Ah, so FA vs PyT is no longer an implementation detail since it expects a different mask format. This makes me think we need an option to explicitly enable or disable FA, and to error out instead of falling back to PyT.
@timmoon10timmoon10Jun 22, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Alternatively, we should expect the mask to be in the format for PyT, and then convert it to the FA format internally so it isn't visible to the user.
@ptrendxptrendxJun 23, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Agree with Tim. OR, accept both in both cases and generate the right one for the current implementation.
To be honest though, the current API for mask type is just bad. For example, what if you want padding AND causal? Then your only option is set the padding mask type and create the causal mask yourself (which would mean you can't do FA in the current implementation which is bad). There is also nothing stopping you really from having arbitrary padding mask (as in, with random elements zeroed out), which would break the assumptions you have in this PR. I think what we could do is introduce new names for the padding type (let's say "pad", "arbitrary" and "no_mask"), deprecate the old names and add causal as separate switch. Then we could say that pad only accepts the "mask" that is actually just list of sequence lengths, arbitrary is whatever you want (and will not go through FA) and causal could be switched irrespective of the mask type.

 if skip_wgrad:
 _disable_wgrads(block)
Expand All@@ -260,26 +247,24 @@ def _test_sanity_e2e_T5(block, bs, dtype, config, fp8_recipe, skip_wgrad):
 te_inp_hidden_states = torch.randn(
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_attn_mask = (
 torch.rand(
 (
 1,
 1,
 config.seq_len,
 config.seq_len,
 )
 )
 .cuda()
 .bool()
 )

 if dtype == torch.float32:
 mask_shape = torch.Size([1, 1, config.seq_len, config.seq_len])
 else:
 mask_shape = torch.Size([config.seq_len, bs])

 enc_dec_attn_mask = torch.rand(mask_shape).cuda().bool()

 if skip_wgrad:
 _disable_wgrads(block)

 use_fp8 = fp8_recipe is not None
 with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
 te_out = block(
 te_inp_hidden_states, te_inp_attn_mask, encoder_output=te_inp_hidden_states
 te_inp_hidden_states,
 None,
 encoder_output=te_inp_hidden_states,
 enc_dec_attn_mask=enc_dec_attn_mask,
 )
 loss = te_out.sum()
 loss.backward()
Expand DownExpand Up@@ -468,12 +453,14 @@ def test_sanity_bert(dtype, bs, fp8_recipe, model, skip_wgrad, zero_centered_gam
 apply_residual_connection_post_layernorm=True,
 output_layernorm=True,
 zero_centered_gamma=zero_centered_gamma,
 self_attn_mask_type="padding",

 )
 .to(dtype=dtype)
 .cuda()
 )

_test_sanity_e2e(block, bs, dtype, config, fp8_recipe, skip_wgrad)
_test_sanity_e2e_bert(block, bs, dtype, config, fp8_recipe, skip_wgrad)


@pytest.mark.parametrize("dtype", param_types)
Expand Down

61 changes: 41 additions & 20 deletions
transformer_engine/pytorch/attention.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -44,6 +44,7 @@
 checkpoint,
)
from transformer_engine.pytorch.export import is_in_onnx_export_mode
from transformer_engine.pytorch.jit import jit_fuser

_flash_attn_version = packaging.version.Version(version("flash-attn"))
_flash_attn_version_required = packaging.version.Version("1.0.6")
Expand All@@ -52,6 +53,19 @@
__all__ = ["DotProductAttention"]


@jit_fuser
def get_cu_seqlens(padding_mask: torch.Tensor) -> torch.Tensor:
 """
 Given a padding mask of shape [seq_len, batch_size], returns an int32
@cyanguwacyanguwaJun 22, 2023

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 guess by seq_len you mean max_seqlen? Also, should the mask tensor be required to be a CUDA tensor?
@ksivamanksivamanJun 22, 2023

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

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

Yeah it is the max_seqlen.
@ksivamanksivamanJun 22, 2023

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

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

Yeah it should be cuda tensor since the inp is cuda as well.
 tensor of shape [batch_size + 1,] containing the cumulative sequence
 lengths of every sample in the batch.
 """
 reduced_mask = padding_mask.sum(dim=0)
 cu_seqlens = reduced_mask.cumsum(dim=0).to(torch.int32)
 zero = torch.zeros(1, dtype=torch.int32, device="cuda")
 return torch.cat((zero, cu_seqlens))


def _rotate_half(x: torch.Tensor) -> torch.Tensor:
 """
 change sign so the last dimension becomes [-odd, +even]
Expand DownExpand Up@@ -345,7 +359,7 @@ def __init__(
 _flash_attn_version >= _flash_attn_version_required
 ), f"FlashAttention minimum version {_flash_attn_version_required} is required."

 self.attn_causal_mask = attn_mask_type == "causal"
 self.attn_mask_type = attn_mask_type
 self.norm_factor = norm_factor
 self.attention_dropout_ctx = attention_dropout_ctx
 self.attention_dropout = attention_dropout
Expand All@@ -356,6 +370,7 @@ def forward(
 query_layer: torch.Tensor,
 key_layer: torch.Tensor,
 value_layer: torch.Tensor,
 attention_mask: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
 """flash-attn fprop"""

Expand All@@ -369,7 +384,6 @@ def forward(
 ), 'FlashAttention currently only supports CUDA tensors.'

 # For now just 128, will make it more general in the future

 if (query_layer.shape[-1] == 128 and
 query_layer.shape[0] * query_layer.shape[1] >= 512 and
 _check_if_interleaved_qkv(query_layer, key_layer, value_layer)):
Expand All@@ -389,18 +403,30 @@ def forward(
 ]

 max_seqlen = seqlen
 cu_seqlens = torch.arange(
 0,
 (batch_size + 1) * seqlen,
 step=seqlen,
 dtype=torch.int32,
 device=query_layer.device)
 if self.attn_mask_type == "padding":
 assert (
 attention_mask is not None
 ), "Boolean attention mask must be provided for padding."
@cyanguwacyanguwaJun 22, 2023

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.

So we're settled on requiring a mask tensor than a cu_seqlen tensor from users?
@ksivamanksivamanJun 22, 2023•

edited

Loading

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

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

Yeah I think so, just to be consistent

 mask_shape = torch.Size([seqlen, batch_size])
 assert (
 attention_mask.shape == mask_shape
 ), f"Expected shape {mask_shape} for attenion mask but found {attention_mask.shape}."

 cu_seqlens = get_cu_seqlens(attention_mask)
 else:
 cu_seqlens = torch.arange(
 0,
 (batch_size + 1) * seqlen,
 step=seqlen,
 dtype=torch.int32,
 device=query_layer.device)

 with self.attention_dropout_ctx():
 output = flash_attn_unpadded_func(
 query_layer, key_layer, value_layer, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen,
 self.attention_dropout if self.training else 0.0,
 softmax_scale=1.0/self.norm_factor, causal=self.attn_causal_mask,
 softmax_scale=1.0/self.norm_factor, causal=self.attn_mask_type=="causal",
 deterministic=self.deterministic,
 )

Expand DownExpand Up@@ -696,7 +722,7 @@ class DotProductAttention(torch.nn.Module):
 .. note::

 Argument :attr:`attention_mask` will be ignored in the `forward` call when
 :attr:`attn_mask_type` is set to `"causal"`.
 :attr:`attn_mask_type` is set to `"causal"` or `"no_mask"`.

 .. warning::

Expand All@@ -714,7 +740,7 @@ class DotProductAttention(torch.nn.Module):
 number of key-value channels.
 attention_dropout: float, default = 0.0
 dropout probability for the dropout op during multi-head attention.
 attn_mask_type: {'causal', 'padding'}, default = `causal`
 attn_mask_type: {'causal', 'padding', 'no_mask'}, default = `causal`
 type of attention mask passed into softmax operation.
 layer_number: int, default = `None`
 layer number of the current `DotProductAttention` when multiple such modules
Expand DownExpand Up@@ -829,7 +855,7 @@ def forward(
 .. note::

 Argument :attr:`attention_mask` will be ignored when :attr:`attn_mask_type`
 is set to `"causal"`.
 is set to `"causal"` or `"no_mask"`.

 .. note::

Expand DownExpand Up@@ -884,7 +910,6 @@ def forward(
 use_flash_attention = False

 if self.attn_mask_type == "padding" and attention_mask is not None:
 use_flash_attention = False
 use_fused_attention = False

 if is_in_onnx_export_mode():
Expand All@@ -911,8 +936,9 @@ def forward(
 return self._checkpointed_attention_forward(self.flash_attention,
 query_layer,
 key_layer,
 value_layer)
 return self.flash_attention(query_layer, key_layer, value_layer)
 value_layer,
 attention_mask)
 return self.flash_attention(query_layer, key_layer, value_layer, attention_mask)

 if use_fused_attention:
 if checkpoint_core_attention:
Expand DownExpand Up@@ -1139,11 +1165,6 @@ def forward(
 """MultiHeadAttention FWD"""
 # hidden_states: [sq, b, h]

 if self.attn_mask_type != "causal" and attention_mask is not None:
 assert (
 attention_mask.dtype == torch.bool
 ), "Attention mask must be a boolean tensor"

 # =================================================
 # Pre-allocate memory for key-values for inference.
 # =================================================
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Auto-enable theater mode on YouTube
(function() {
function tryTheater() {
var btn = document.querySelector('button[aria-label="Theater mode"], ytd-player #player button[title="Theater mode"]');
if (btn && !btn.classList.contains('activated')) {
btn.click();
}
}
// Try immediately
tryTheater();
// Try after navigation (SPA)
var lastUrl = location.href;
setInterval(function() {
if (location.href !== lastUrl) {
lastUrl = location.href;
setTimeout(tryTheater, 500);
}
}, 1000);
// Also try on player load
var observer = new MutationObserver(tryTheater);
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
Expand FA coverage to padding masks by ksivaman · Pull Request #291 · NVIDIA/TransformerEngine · GitHub
Skip to content
Expand FA coverage to padding masks#291

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

 Closed

ksivaman
wants to merge
12
commits into

from

 Closed


Expand FA coverage to padding masks
#291
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


12 commits


Select commit
Hold shift + click to select a range
9a28af3

initial changes [wip]


ksivaman Jun 21, 2023
7b01a27

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 21, 2023
a89e43c

Add padding mask support for FA


ksivaman Jun 21, 2023
ae104b3

Address review comments


ksivaman Jun 21, 2023
e38d18d

rm causal mask from tests and add padding


ksivaman Jun 22, 2023
b0212bd

fix some conflicts


ksivaman Jun 22, 2023
e3dfd8b

conflicts


ksivaman Jun 22, 2023
ad8fc31

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 22, 2023
e2b2ddf

Fixes


ksivaman Jun 22, 2023
8eb8d90

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 23, 2023
e0cd408

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 26, 2023
e181188

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 27, 2023

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


13 changes: 5 additions & 8 deletions
tests/pytorch/test_numerics.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -401,11 +401,10 @@ def _test_e2e_selective_recompute(block, bs, dtype, config, recompute=False):
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_hidden_states.retain_grad()
 te_inp_attn_mask = get_causal_attn_mask(config.seq_len)

 te_out = block(
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 checkpoint_core_attention=recompute,
 )
 loss = te_out.sum()
Expand DownExpand Up@@ -461,7 +460,6 @@ def _test_e2e_full_recompute(block, bs, dtype, config, recompute=False):
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_hidden_states.retain_grad()
 te_inp_attn_mask = get_causal_attn_mask(config.seq_len)

 if recompute:
 te_out = te_checkpoint(
Expand All@@ -470,13 +468,13 @@ def _test_e2e_full_recompute(block, bs, dtype, config, recompute=False):
 get_dummy_cuda_rng_tracker,
 None, # tp_group
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 checkpoint_core_attention=False,
 )
 else:
 te_out = block(
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 checkpoint_core_attention=False,
 )
 loss = te_out.sum()
Expand DownExpand Up@@ -556,14 +554,13 @@ def _test_e2e_checkpointing(bs, dtype, config, checkpoint=False, steps=10, path=
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_hidden_states.retain_grad()
 te_inp_attn_mask = get_causal_attn_mask(config.seq_len)

 block = _test_e2e_checkpointing_get_model(config, dtype)

 for _ in range(steps // 2):
 te_out = block(
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 )
 loss = te_out.sum()
 loss.backward()
Expand DownExpand Up@@ -594,7 +591,7 @@ def _test_e2e_checkpointing(bs, dtype, config, checkpoint=False, steps=10, path=
 for _ in range(steps // 2):
 te_out = block(
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 )
 loss = te_out.sum()
 loss.backward()
Expand Down

91 changes: 39 additions & 52 deletions
tests/pytorch/test_sanity.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -154,26 +154,14 @@ def _test_sanity_e2e_amp(block, bs, dtype, config, fp8_recipe, skip_wgrad):
 config.seq_len, bs, config.hidden_size, dtype=torch.float32, requires_grad=True
 ).cuda()
 te_inp_hidden_states.retain_grad()
 te_inp_attn_mask = (
 torch.rand(
 (
 1,
 1,
 config.seq_len,
 config.seq_len,
 )
 )
 .cuda()
 .bool()
 )

 if skip_wgrad:
 _disable_wgrads(block)

 use_fp8 = fp8_recipe is not None
 with torch.autocast(device_type="cuda", enabled=True, dtype=dtype):
 with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
 te_out = block(te_inp_hidden_states, te_inp_attn_mask)
 te_out = block(te_inp_hidden_states)
 loss = te_out.sum()

 loss.backward()
Expand All@@ -190,18 +178,6 @@ def _test_sanity_e2e_gradient_accumulation_fusion(block, bs, dtype, config, fp8_
 te_inp_hidden_states = torch.randn(
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_attn_mask = (
 torch.rand(
 (
 1,
 1,
 config.seq_len,
 config.seq_len,
 )
 )
 .cuda()
 .bool()
 )

 if skip_wgrad:
 _disable_wgrads(block)
Expand All@@ -214,7 +190,7 @@ def _test_sanity_e2e_gradient_accumulation_fusion(block, bs, dtype, config, fp8_

 use_fp8 = fp8_recipe is not None
 with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
 te_out = block(te_inp_hidden_states, te_inp_attn_mask)
 te_out = block(te_inp_hidden_states)
 loss = te_out.sum()
 loss.backward()
 torch.cuda.synchronize()
Expand All@@ -232,18 +208,29 @@ def _test_sanity_e2e(block, bs, dtype, config, fp8_recipe, skip_wgrad):
 te_inp_hidden_states = torch.randn(
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_attn_mask = (
 torch.rand(
 (
 1,
 1,
 config.seq_len,
 config.seq_len,
 )
 )
 .cuda()
 .bool()
 )

 if skip_wgrad:
 _disable_wgrads(block)

 use_fp8 = fp8_recipe is not None
 with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
 te_out = block(te_inp_hidden_states)
 loss = te_out.sum()
 loss.backward()
 torch.cuda.synchronize()


def _test_sanity_e2e_bert(block, bs, dtype, config, fp8_recipe, skip_wgrad):
 te_inp_hidden_states = torch.randn(
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()

 if dtype == torch.float32:
 mask_shape = torch.Size([1, 1, config.seq_len, config.seq_len])
 else:
 mask_shape = torch.Size([config.seq_len, bs])

 te_inp_attn_mask = torch.rand(mask_shape).cuda().bool()

Comment on lines

+228
 to 
+233

@timmoon10timmoon10Jun 21, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Why does the FP32 test require a 4D mask while other dtypes use a 2D mask?
@ksivamanksivamanJun 21, 2023

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

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

The FP32 doesn't take the FA path and so uses the PyTorch torch softmax path for padding mask which expects this
@timmoon10timmoon10Jun 22, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Ah, so FA vs PyT is no longer an implementation detail since it expects a different mask format. This makes me think we need an option to explicitly enable or disable FA, and to error out instead of falling back to PyT.
@timmoon10timmoon10Jun 22, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Alternatively, we should expect the mask to be in the format for PyT, and then convert it to the FA format internally so it isn't visible to the user.
@ptrendxptrendxJun 23, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Agree with Tim. OR, accept both in both cases and generate the right one for the current implementation.
To be honest though, the current API for mask type is just bad. For example, what if you want padding AND causal? Then your only option is set the padding mask type and create the causal mask yourself (which would mean you can't do FA in the current implementation which is bad). There is also nothing stopping you really from having arbitrary padding mask (as in, with random elements zeroed out), which would break the assumptions you have in this PR. I think what we could do is introduce new names for the padding type (let's say "pad", "arbitrary" and "no_mask"), deprecate the old names and add causal as separate switch. Then we could say that pad only accepts the "mask" that is actually just list of sequence lengths, arbitrary is whatever you want (and will not go through FA) and causal could be switched irrespective of the mask type.

 if skip_wgrad:
 _disable_wgrads(block)
Expand All@@ -260,26 +247,24 @@ def _test_sanity_e2e_T5(block, bs, dtype, config, fp8_recipe, skip_wgrad):
 te_inp_hidden_states = torch.randn(
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_attn_mask = (
 torch.rand(
 (
 1,
 1,
 config.seq_len,
 config.seq_len,
 )
 )
 .cuda()
 .bool()
 )

 if dtype == torch.float32:
 mask_shape = torch.Size([1, 1, config.seq_len, config.seq_len])
 else:
 mask_shape = torch.Size([config.seq_len, bs])

 enc_dec_attn_mask = torch.rand(mask_shape).cuda().bool()

 if skip_wgrad:
 _disable_wgrads(block)

 use_fp8 = fp8_recipe is not None
 with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
 te_out = block(
 te_inp_hidden_states, te_inp_attn_mask, encoder_output=te_inp_hidden_states
 te_inp_hidden_states,
 None,
 encoder_output=te_inp_hidden_states,
 enc_dec_attn_mask=enc_dec_attn_mask,
 )
 loss = te_out.sum()
 loss.backward()
Expand DownExpand Up@@ -468,12 +453,14 @@ def test_sanity_bert(dtype, bs, fp8_recipe, model, skip_wgrad, zero_centered_gam
 apply_residual_connection_post_layernorm=True,
 output_layernorm=True,
 zero_centered_gamma=zero_centered_gamma,
 self_attn_mask_type="padding",

 )
 .to(dtype=dtype)
 .cuda()
 )

_test_sanity_e2e(block, bs, dtype, config, fp8_recipe, skip_wgrad)
_test_sanity_e2e_bert(block, bs, dtype, config, fp8_recipe, skip_wgrad)


@pytest.mark.parametrize("dtype", param_types)
Expand Down

61 changes: 41 additions & 20 deletions
transformer_engine/pytorch/attention.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -44,6 +44,7 @@
 checkpoint,
)
from transformer_engine.pytorch.export import is_in_onnx_export_mode
from transformer_engine.pytorch.jit import jit_fuser

_flash_attn_version = packaging.version.Version(version("flash-attn"))
_flash_attn_version_required = packaging.version.Version("1.0.6")
Expand All@@ -52,6 +53,19 @@
__all__ = ["DotProductAttention"]


@jit_fuser
def get_cu_seqlens(padding_mask: torch.Tensor) -> torch.Tensor:
 """
 Given a padding mask of shape [seq_len, batch_size], returns an int32
@cyanguwacyanguwaJun 22, 2023

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 guess by seq_len you mean max_seqlen? Also, should the mask tensor be required to be a CUDA tensor?
@ksivamanksivamanJun 22, 2023

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

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

Yeah it is the max_seqlen.
@ksivamanksivamanJun 22, 2023

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

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

Yeah it should be cuda tensor since the inp is cuda as well.
 tensor of shape [batch_size + 1,] containing the cumulative sequence
 lengths of every sample in the batch.
 """
 reduced_mask = padding_mask.sum(dim=0)
 cu_seqlens = reduced_mask.cumsum(dim=0).to(torch.int32)
 zero = torch.zeros(1, dtype=torch.int32, device="cuda")
 return torch.cat((zero, cu_seqlens))


def _rotate_half(x: torch.Tensor) -> torch.Tensor:
 """
 change sign so the last dimension becomes [-odd, +even]
Expand DownExpand Up@@ -345,7 +359,7 @@ def __init__(
 _flash_attn_version >= _flash_attn_version_required
 ), f"FlashAttention minimum version {_flash_attn_version_required} is required."

 self.attn_causal_mask = attn_mask_type == "causal"
 self.attn_mask_type = attn_mask_type
 self.norm_factor = norm_factor
 self.attention_dropout_ctx = attention_dropout_ctx
 self.attention_dropout = attention_dropout
Expand All@@ -356,6 +370,7 @@ def forward(
 query_layer: torch.Tensor,
 key_layer: torch.Tensor,
 value_layer: torch.Tensor,
 attention_mask: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
 """flash-attn fprop"""

Expand All@@ -369,7 +384,6 @@ def forward(
 ), 'FlashAttention currently only supports CUDA tensors.'

 # For now just 128, will make it more general in the future

 if (query_layer.shape[-1] == 128 and
 query_layer.shape[0] * query_layer.shape[1] >= 512 and
 _check_if_interleaved_qkv(query_layer, key_layer, value_layer)):
Expand All@@ -389,18 +403,30 @@ def forward(
 ]

 max_seqlen = seqlen
 cu_seqlens = torch.arange(
 0,
 (batch_size + 1) * seqlen,
 step=seqlen,
 dtype=torch.int32,
 device=query_layer.device)
 if self.attn_mask_type == "padding":
 assert (
 attention_mask is not None
 ), "Boolean attention mask must be provided for padding."
@cyanguwacyanguwaJun 22, 2023

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.

So we're settled on requiring a mask tensor than a cu_seqlen tensor from users?
@ksivamanksivamanJun 22, 2023•

edited

Loading

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

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

Yeah I think so, just to be consistent

 mask_shape = torch.Size([seqlen, batch_size])
 assert (
 attention_mask.shape == mask_shape
 ), f"Expected shape {mask_shape} for attenion mask but found {attention_mask.shape}."

 cu_seqlens = get_cu_seqlens(attention_mask)
 else:
 cu_seqlens = torch.arange(
 0,
 (batch_size + 1) * seqlen,
 step=seqlen,
 dtype=torch.int32,
 device=query_layer.device)

 with self.attention_dropout_ctx():
 output = flash_attn_unpadded_func(
 query_layer, key_layer, value_layer, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen,
 self.attention_dropout if self.training else 0.0,
 softmax_scale=1.0/self.norm_factor, causal=self.attn_causal_mask,
 softmax_scale=1.0/self.norm_factor, causal=self.attn_mask_type=="causal",
 deterministic=self.deterministic,
 )

Expand DownExpand Up@@ -696,7 +722,7 @@ class DotProductAttention(torch.nn.Module):
 .. note::

 Argument :attr:`attention_mask` will be ignored in the `forward` call when
 :attr:`attn_mask_type` is set to `"causal"`.
 :attr:`attn_mask_type` is set to `"causal"` or `"no_mask"`.

 .. warning::

Expand All@@ -714,7 +740,7 @@ class DotProductAttention(torch.nn.Module):
 number of key-value channels.
 attention_dropout: float, default = 0.0
 dropout probability for the dropout op during multi-head attention.
 attn_mask_type: {'causal', 'padding'}, default = `causal`
 attn_mask_type: {'causal', 'padding', 'no_mask'}, default = `causal`
 type of attention mask passed into softmax operation.
 layer_number: int, default = `None`
 layer number of the current `DotProductAttention` when multiple such modules
Expand DownExpand Up@@ -829,7 +855,7 @@ def forward(
 .. note::

 Argument :attr:`attention_mask` will be ignored when :attr:`attn_mask_type`
 is set to `"causal"`.
 is set to `"causal"` or `"no_mask"`.

 .. note::

Expand DownExpand Up@@ -884,7 +910,6 @@ def forward(
 use_flash_attention = False

 if self.attn_mask_type == "padding" and attention_mask is not None:
 use_flash_attention = False
 use_fused_attention = False

 if is_in_onnx_export_mode():
Expand All@@ -911,8 +936,9 @@ def forward(
 return self._checkpointed_attention_forward(self.flash_attention,
 query_layer,
 key_layer,
 value_layer)
 return self.flash_attention(query_layer, key_layer, value_layer)
 value_layer,
 attention_mask)
 return self.flash_attention(query_layer, key_layer, value_layer, attention_mask)

 if use_fused_attention:
 if checkpoint_core_attention:
Expand DownExpand Up@@ -1139,11 +1165,6 @@ def forward(
 """MultiHeadAttention FWD"""
 # hidden_states: [sq, b, h]

 if self.attn_mask_type != "causal" and attention_mask is not None:
 assert (
 attention_mask.dtype == torch.bool
 ), "Attention mask must be a boolean tensor"

 # =================================================
 # Pre-allocate memory for key-values for inference.
 # =================================================
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Remove or un-stick sticky/fixed headers that block content
(function() {
function unstick() {
document.querySelectorAll('header, nav, [role="banner"], .header, .navbar, .sticky, .fixed-top, [style*="position: fixed"], [style*="position:sticky"]').forEach(function(el) {
if (el.style.position === 'fixed' || el.style.position === 'sticky' || getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {
el.style.position = 'static';
el.style.top = 'auto';
el.style.zIndex = 'auto';
}
});
}
unstick();
var observer = new MutationObserver(unstick);
observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });
})();
}
} catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
Expand FA coverage to padding masks by ksivaman · Pull Request #291 · NVIDIA/TransformerEngine · GitHub
Skip to content
Expand FA coverage to padding masks#291

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

 Closed

ksivaman
wants to merge
12
commits into

from

 Closed


Expand FA coverage to padding masks
#291
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


12 commits


Select commit
Hold shift + click to select a range
9a28af3

initial changes [wip]


ksivaman Jun 21, 2023
7b01a27

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 21, 2023
a89e43c

Add padding mask support for FA


ksivaman Jun 21, 2023
ae104b3

Address review comments


ksivaman Jun 21, 2023
e38d18d

rm causal mask from tests and add padding


ksivaman Jun 22, 2023
b0212bd

fix some conflicts


ksivaman Jun 22, 2023
e3dfd8b

conflicts


ksivaman Jun 22, 2023
ad8fc31

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 22, 2023
e2b2ddf

Fixes


ksivaman Jun 22, 2023
8eb8d90

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 23, 2023
e0cd408

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 26, 2023
e181188

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 27, 2023

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


13 changes: 5 additions & 8 deletions
tests/pytorch/test_numerics.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -401,11 +401,10 @@ def _test_e2e_selective_recompute(block, bs, dtype, config, recompute=False):
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_hidden_states.retain_grad()
 te_inp_attn_mask = get_causal_attn_mask(config.seq_len)

 te_out = block(
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 checkpoint_core_attention=recompute,
 )
 loss = te_out.sum()
Expand DownExpand Up@@ -461,7 +460,6 @@ def _test_e2e_full_recompute(block, bs, dtype, config, recompute=False):
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_hidden_states.retain_grad()
 te_inp_attn_mask = get_causal_attn_mask(config.seq_len)

 if recompute:
 te_out = te_checkpoint(
Expand All@@ -470,13 +468,13 @@ def _test_e2e_full_recompute(block, bs, dtype, config, recompute=False):
 get_dummy_cuda_rng_tracker,
 None, # tp_group
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 checkpoint_core_attention=False,
 )
 else:
 te_out = block(
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 checkpoint_core_attention=False,
 )
 loss = te_out.sum()
Expand DownExpand Up@@ -556,14 +554,13 @@ def _test_e2e_checkpointing(bs, dtype, config, checkpoint=False, steps=10, path=
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_hidden_states.retain_grad()
 te_inp_attn_mask = get_causal_attn_mask(config.seq_len)

 block = _test_e2e_checkpointing_get_model(config, dtype)

 for _ in range(steps // 2):
 te_out = block(
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 )
 loss = te_out.sum()
 loss.backward()
Expand DownExpand Up@@ -594,7 +591,7 @@ def _test_e2e_checkpointing(bs, dtype, config, checkpoint=False, steps=10, path=
 for _ in range(steps // 2):
 te_out = block(
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 )
 loss = te_out.sum()
 loss.backward()
Expand Down

91 changes: 39 additions & 52 deletions
tests/pytorch/test_sanity.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -154,26 +154,14 @@ def _test_sanity_e2e_amp(block, bs, dtype, config, fp8_recipe, skip_wgrad):
 config.seq_len, bs, config.hidden_size, dtype=torch.float32, requires_grad=True
 ).cuda()
 te_inp_hidden_states.retain_grad()
 te_inp_attn_mask = (
 torch.rand(
 (
 1,
 1,
 config.seq_len,
 config.seq_len,
 )
 )
 .cuda()
 .bool()
 )

 if skip_wgrad:
 _disable_wgrads(block)

 use_fp8 = fp8_recipe is not None
 with torch.autocast(device_type="cuda", enabled=True, dtype=dtype):
 with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
 te_out = block(te_inp_hidden_states, te_inp_attn_mask)
 te_out = block(te_inp_hidden_states)
 loss = te_out.sum()

 loss.backward()
Expand All@@ -190,18 +178,6 @@ def _test_sanity_e2e_gradient_accumulation_fusion(block, bs, dtype, config, fp8_
 te_inp_hidden_states = torch.randn(
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_attn_mask = (
 torch.rand(
 (
 1,
 1,
 config.seq_len,
 config.seq_len,
 )
 )
 .cuda()
 .bool()
 )

 if skip_wgrad:
 _disable_wgrads(block)
Expand All@@ -214,7 +190,7 @@ def _test_sanity_e2e_gradient_accumulation_fusion(block, bs, dtype, config, fp8_

 use_fp8 = fp8_recipe is not None
 with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
 te_out = block(te_inp_hidden_states, te_inp_attn_mask)
 te_out = block(te_inp_hidden_states)
 loss = te_out.sum()
 loss.backward()
 torch.cuda.synchronize()
Expand All@@ -232,18 +208,29 @@ def _test_sanity_e2e(block, bs, dtype, config, fp8_recipe, skip_wgrad):
 te_inp_hidden_states = torch.randn(
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_attn_mask = (
 torch.rand(
 (
 1,
 1,
 config.seq_len,
 config.seq_len,
 )
 )
 .cuda()
 .bool()
 )

 if skip_wgrad:
 _disable_wgrads(block)

 use_fp8 = fp8_recipe is not None
 with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
 te_out = block(te_inp_hidden_states)
 loss = te_out.sum()
 loss.backward()
 torch.cuda.synchronize()


def _test_sanity_e2e_bert(block, bs, dtype, config, fp8_recipe, skip_wgrad):
 te_inp_hidden_states = torch.randn(
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()

 if dtype == torch.float32:
 mask_shape = torch.Size([1, 1, config.seq_len, config.seq_len])
 else:
 mask_shape = torch.Size([config.seq_len, bs])

 te_inp_attn_mask = torch.rand(mask_shape).cuda().bool()

Comment on lines

+228
 to 
+233

@timmoon10timmoon10Jun 21, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Why does the FP32 test require a 4D mask while other dtypes use a 2D mask?
@ksivamanksivamanJun 21, 2023

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

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

The FP32 doesn't take the FA path and so uses the PyTorch torch softmax path for padding mask which expects this
@timmoon10timmoon10Jun 22, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Ah, so FA vs PyT is no longer an implementation detail since it expects a different mask format. This makes me think we need an option to explicitly enable or disable FA, and to error out instead of falling back to PyT.
@timmoon10timmoon10Jun 22, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Alternatively, we should expect the mask to be in the format for PyT, and then convert it to the FA format internally so it isn't visible to the user.
@ptrendxptrendxJun 23, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Agree with Tim. OR, accept both in both cases and generate the right one for the current implementation.
To be honest though, the current API for mask type is just bad. For example, what if you want padding AND causal? Then your only option is set the padding mask type and create the causal mask yourself (which would mean you can't do FA in the current implementation which is bad). There is also nothing stopping you really from having arbitrary padding mask (as in, with random elements zeroed out), which would break the assumptions you have in this PR. I think what we could do is introduce new names for the padding type (let's say "pad", "arbitrary" and "no_mask"), deprecate the old names and add causal as separate switch. Then we could say that pad only accepts the "mask" that is actually just list of sequence lengths, arbitrary is whatever you want (and will not go through FA) and causal could be switched irrespective of the mask type.

 if skip_wgrad:
 _disable_wgrads(block)
Expand All@@ -260,26 +247,24 @@ def _test_sanity_e2e_T5(block, bs, dtype, config, fp8_recipe, skip_wgrad):
 te_inp_hidden_states = torch.randn(
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_attn_mask = (
 torch.rand(
 (
 1,
 1,
 config.seq_len,
 config.seq_len,
 )
 )
 .cuda()
 .bool()
 )

 if dtype == torch.float32:
 mask_shape = torch.Size([1, 1, config.seq_len, config.seq_len])
 else:
 mask_shape = torch.Size([config.seq_len, bs])

 enc_dec_attn_mask = torch.rand(mask_shape).cuda().bool()

 if skip_wgrad:
 _disable_wgrads(block)

 use_fp8 = fp8_recipe is not None
 with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
 te_out = block(
 te_inp_hidden_states, te_inp_attn_mask, encoder_output=te_inp_hidden_states
 te_inp_hidden_states,
 None,
 encoder_output=te_inp_hidden_states,
 enc_dec_attn_mask=enc_dec_attn_mask,
 )
 loss = te_out.sum()
 loss.backward()
Expand DownExpand Up@@ -468,12 +453,14 @@ def test_sanity_bert(dtype, bs, fp8_recipe, model, skip_wgrad, zero_centered_gam
 apply_residual_connection_post_layernorm=True,
 output_layernorm=True,
 zero_centered_gamma=zero_centered_gamma,
 self_attn_mask_type="padding",

 )
 .to(dtype=dtype)
 .cuda()
 )

_test_sanity_e2e(block, bs, dtype, config, fp8_recipe, skip_wgrad)
_test_sanity_e2e_bert(block, bs, dtype, config, fp8_recipe, skip_wgrad)


@pytest.mark.parametrize("dtype", param_types)
Expand Down

61 changes: 41 additions & 20 deletions
transformer_engine/pytorch/attention.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -44,6 +44,7 @@
 checkpoint,
)
from transformer_engine.pytorch.export import is_in_onnx_export_mode
from transformer_engine.pytorch.jit import jit_fuser

_flash_attn_version = packaging.version.Version(version("flash-attn"))
_flash_attn_version_required = packaging.version.Version("1.0.6")
Expand All@@ -52,6 +53,19 @@
__all__ = ["DotProductAttention"]


@jit_fuser
def get_cu_seqlens(padding_mask: torch.Tensor) -> torch.Tensor:
 """
 Given a padding mask of shape [seq_len, batch_size], returns an int32
@cyanguwacyanguwaJun 22, 2023

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 guess by seq_len you mean max_seqlen? Also, should the mask tensor be required to be a CUDA tensor?
@ksivamanksivamanJun 22, 2023

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

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

Yeah it is the max_seqlen.
@ksivamanksivamanJun 22, 2023

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

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

Yeah it should be cuda tensor since the inp is cuda as well.
 tensor of shape [batch_size + 1,] containing the cumulative sequence
 lengths of every sample in the batch.
 """
 reduced_mask = padding_mask.sum(dim=0)
 cu_seqlens = reduced_mask.cumsum(dim=0).to(torch.int32)
 zero = torch.zeros(1, dtype=torch.int32, device="cuda")
 return torch.cat((zero, cu_seqlens))


def _rotate_half(x: torch.Tensor) -> torch.Tensor:
 """
 change sign so the last dimension becomes [-odd, +even]
Expand DownExpand Up@@ -345,7 +359,7 @@ def __init__(
 _flash_attn_version >= _flash_attn_version_required
 ), f"FlashAttention minimum version {_flash_attn_version_required} is required."

 self.attn_causal_mask = attn_mask_type == "causal"
 self.attn_mask_type = attn_mask_type
 self.norm_factor = norm_factor
 self.attention_dropout_ctx = attention_dropout_ctx
 self.attention_dropout = attention_dropout
Expand All@@ -356,6 +370,7 @@ def forward(
 query_layer: torch.Tensor,
 key_layer: torch.Tensor,
 value_layer: torch.Tensor,
 attention_mask: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
 """flash-attn fprop"""

Expand All@@ -369,7 +384,6 @@ def forward(
 ), 'FlashAttention currently only supports CUDA tensors.'

 # For now just 128, will make it more general in the future

 if (query_layer.shape[-1] == 128 and
 query_layer.shape[0] * query_layer.shape[1] >= 512 and
 _check_if_interleaved_qkv(query_layer, key_layer, value_layer)):
Expand All@@ -389,18 +403,30 @@ def forward(
 ]

 max_seqlen = seqlen
 cu_seqlens = torch.arange(
 0,
 (batch_size + 1) * seqlen,
 step=seqlen,
 dtype=torch.int32,
 device=query_layer.device)
 if self.attn_mask_type == "padding":
 assert (
 attention_mask is not None
 ), "Boolean attention mask must be provided for padding."
@cyanguwacyanguwaJun 22, 2023

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.

So we're settled on requiring a mask tensor than a cu_seqlen tensor from users?
@ksivamanksivamanJun 22, 2023•

edited

Loading

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

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

Yeah I think so, just to be consistent

 mask_shape = torch.Size([seqlen, batch_size])
 assert (
 attention_mask.shape == mask_shape
 ), f"Expected shape {mask_shape} for attenion mask but found {attention_mask.shape}."

 cu_seqlens = get_cu_seqlens(attention_mask)
 else:
 cu_seqlens = torch.arange(
 0,
 (batch_size + 1) * seqlen,
 step=seqlen,
 dtype=torch.int32,
 device=query_layer.device)

 with self.attention_dropout_ctx():
 output = flash_attn_unpadded_func(
 query_layer, key_layer, value_layer, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen,
 self.attention_dropout if self.training else 0.0,
 softmax_scale=1.0/self.norm_factor, causal=self.attn_causal_mask,
 softmax_scale=1.0/self.norm_factor, causal=self.attn_mask_type=="causal",
 deterministic=self.deterministic,
 )

Expand DownExpand Up@@ -696,7 +722,7 @@ class DotProductAttention(torch.nn.Module):
 .. note::

 Argument :attr:`attention_mask` will be ignored in the `forward` call when
 :attr:`attn_mask_type` is set to `"causal"`.
 :attr:`attn_mask_type` is set to `"causal"` or `"no_mask"`.

 .. warning::

Expand All@@ -714,7 +740,7 @@ class DotProductAttention(torch.nn.Module):
 number of key-value channels.
 attention_dropout: float, default = 0.0
 dropout probability for the dropout op during multi-head attention.
 attn_mask_type: {'causal', 'padding'}, default = `causal`
 attn_mask_type: {'causal', 'padding', 'no_mask'}, default = `causal`
 type of attention mask passed into softmax operation.
 layer_number: int, default = `None`
 layer number of the current `DotProductAttention` when multiple such modules
Expand DownExpand Up@@ -829,7 +855,7 @@ def forward(
 .. note::

 Argument :attr:`attention_mask` will be ignored when :attr:`attn_mask_type`
 is set to `"causal"`.
 is set to `"causal"` or `"no_mask"`.

 .. note::

Expand DownExpand Up@@ -884,7 +910,6 @@ def forward(
 use_flash_attention = False

 if self.attn_mask_type == "padding" and attention_mask is not None:
 use_flash_attention = False
 use_fused_attention = False

 if is_in_onnx_export_mode():
Expand All@@ -911,8 +936,9 @@ def forward(
 return self._checkpointed_attention_forward(self.flash_attention,
 query_layer,
 key_layer,
 value_layer)
 return self.flash_attention(query_layer, key_layer, value_layer)
 value_layer,
 attention_mask)
 return self.flash_attention(query_layer, key_layer, value_layer, attention_mask)

 if use_fused_attention:
 if checkpoint_core_attention:
Expand DownExpand Up@@ -1139,11 +1165,6 @@ def forward(
 """MultiHeadAttention FWD"""
 # hidden_states: [sq, b, h]

 if self.attn_mask_type != "causal" and attention_mask is not None:
 assert (
 attention_mask.dtype == torch.bool
 ), "Attention mask must be a boolean tensor"

 # =================================================
 # Pre-allocate memory for key-values for inference.
 # =================================================
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Universal Dark Mode - works on any site
(function() {
var enabled = true;
function applyDarkMode() {
if (!enabled) return;
// Create style element if it doesn't exist
var style = document.getElementById('universal-dark-mode-style');
if (!style) {
style = document.createElement('style');
style.id = 'universal-dark-mode-style';
document.head.appendChild(style);
}
// Dark mode CSS - inverts colors but preserves images/video
style.textContent = '
/* Invert everything except media */
html {
filter: invert(1) hue-rotate(180deg) !important;
background: #1a1a2e !important;
}
/* Restore images, videos, iframes, canvas */
img, video, iframe, canvas, svg, picture, [style*="background-image"] {
filter: invert(1) hue-rotate(180deg) !important;
}
/* Preserve specific elements that should not be inverted */
.no-dark-mode, .no-dark-mode *,
[data-theme="light"], [data-theme="light"],
.ace_editor, .ace_editor *,
.CodeMirror, .CodeMirror *,
.monaco-editor, .monaco-editor *,
.markdown-body pre, .markdown-body pre *,
.highlight, .highlight *,
pre code, pre code * {
filter: none !important;
}
/* Fix common UI elements */
.modal, .popup, .dropdown-menu, .tooltip, .popover {
filter: invert(1) hue-rotate(180deg) !important;
background: #2d2d44 !important;
border-color: #444 !important;
}
/* Scrollbars */
::-webkit-scrollbar { background: #1a1a2e !important; }
::-webkit-scrollbar-thumb { background: #444 !important; }
::-webkit-scrollbar-thumb:hover { background: #555 !important; }
/* Selection */
::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }
::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }
';
}
function removeDarkMode() {
var style = document.getElementById('universal-dark-mode-style');
if (style) style.remove();
}
// Toggle with Alt+Shift+D
document.addEventListener('keydown', function(e) {
if (e.altKey && e.shiftKey && e.key === 'D') {
e.preventDefault();
enabled = !enabled;
if (enabled) {
applyDarkMode();
console.log('[Universal Dark Mode] Enabled');
} else {
removeDarkMode();
console.log('[Universal Dark Mode] Disabled');
}
}
});
// Apply on load
applyDarkMode();
// Re-apply on dynamic content
var observer = new MutationObserver(function(mutations) {
if (enabled && !document.getElementById('universal-dark-mode-style')) {
applyDarkMode();
}
});
observer.observe(document.head, { childList: true });
console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');
})();
}
} catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); }
})();
})();
Expand FA coverage to padding masks by ksivaman · Pull Request #291 · NVIDIA/TransformerEngine · GitHub
Skip to content
Expand FA coverage to padding masks#291

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

 Closed

ksivaman
wants to merge
12
commits into

from

 Closed


Expand FA coverage to padding masks
#291
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


12 commits


Select commit
Hold shift + click to select a range
9a28af3

initial changes [wip]


ksivaman Jun 21, 2023
7b01a27

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 21, 2023
a89e43c

Add padding mask support for FA


ksivaman Jun 21, 2023
ae104b3

Address review comments


ksivaman Jun 21, 2023
e38d18d

rm causal mask from tests and add padding


ksivaman Jun 22, 2023
b0212bd

fix some conflicts


ksivaman Jun 22, 2023
e3dfd8b

conflicts


ksivaman Jun 22, 2023
ad8fc31

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 22, 2023
e2b2ddf

Fixes


ksivaman Jun 22, 2023
8eb8d90

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 23, 2023
e0cd408

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 26, 2023
e181188

Merge branch 'main' into expand_flash_attn_coverage


ksivaman Jun 27, 2023

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


13 changes: 5 additions & 8 deletions
tests/pytorch/test_numerics.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -401,11 +401,10 @@ def _test_e2e_selective_recompute(block, bs, dtype, config, recompute=False):
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_hidden_states.retain_grad()
 te_inp_attn_mask = get_causal_attn_mask(config.seq_len)

 te_out = block(
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 checkpoint_core_attention=recompute,
 )
 loss = te_out.sum()
Expand DownExpand Up@@ -461,7 +460,6 @@ def _test_e2e_full_recompute(block, bs, dtype, config, recompute=False):
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_hidden_states.retain_grad()
 te_inp_attn_mask = get_causal_attn_mask(config.seq_len)

 if recompute:
 te_out = te_checkpoint(
Expand All@@ -470,13 +468,13 @@ def _test_e2e_full_recompute(block, bs, dtype, config, recompute=False):
 get_dummy_cuda_rng_tracker,
 None, # tp_group
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 checkpoint_core_attention=False,
 )
 else:
 te_out = block(
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 checkpoint_core_attention=False,
 )
 loss = te_out.sum()
Expand DownExpand Up@@ -556,14 +554,13 @@ def _test_e2e_checkpointing(bs, dtype, config, checkpoint=False, steps=10, path=
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_hidden_states.retain_grad()
 te_inp_attn_mask = get_causal_attn_mask(config.seq_len)

 block = _test_e2e_checkpointing_get_model(config, dtype)

 for _ in range(steps // 2):
 te_out = block(
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 )
 loss = te_out.sum()
 loss.backward()
Expand DownExpand Up@@ -594,7 +591,7 @@ def _test_e2e_checkpointing(bs, dtype, config, checkpoint=False, steps=10, path=
 for _ in range(steps // 2):
 te_out = block(
 te_inp_hidden_states,
te_inp_attn_mask,
None,
 )
 loss = te_out.sum()
 loss.backward()
Expand Down

91 changes: 39 additions & 52 deletions
tests/pytorch/test_sanity.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -154,26 +154,14 @@ def _test_sanity_e2e_amp(block, bs, dtype, config, fp8_recipe, skip_wgrad):
 config.seq_len, bs, config.hidden_size, dtype=torch.float32, requires_grad=True
 ).cuda()
 te_inp_hidden_states.retain_grad()
 te_inp_attn_mask = (
 torch.rand(
 (
 1,
 1,
 config.seq_len,
 config.seq_len,
 )
 )
 .cuda()
 .bool()
 )

 if skip_wgrad:
 _disable_wgrads(block)

 use_fp8 = fp8_recipe is not None
 with torch.autocast(device_type="cuda", enabled=True, dtype=dtype):
 with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
 te_out = block(te_inp_hidden_states, te_inp_attn_mask)
 te_out = block(te_inp_hidden_states)
 loss = te_out.sum()

 loss.backward()
Expand All@@ -190,18 +178,6 @@ def _test_sanity_e2e_gradient_accumulation_fusion(block, bs, dtype, config, fp8_
 te_inp_hidden_states = torch.randn(
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_attn_mask = (
 torch.rand(
 (
 1,
 1,
 config.seq_len,
 config.seq_len,
 )
 )
 .cuda()
 .bool()
 )

 if skip_wgrad:
 _disable_wgrads(block)
Expand All@@ -214,7 +190,7 @@ def _test_sanity_e2e_gradient_accumulation_fusion(block, bs, dtype, config, fp8_

 use_fp8 = fp8_recipe is not None
 with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
 te_out = block(te_inp_hidden_states, te_inp_attn_mask)
 te_out = block(te_inp_hidden_states)
 loss = te_out.sum()
 loss.backward()
 torch.cuda.synchronize()
Expand All@@ -232,18 +208,29 @@ def _test_sanity_e2e(block, bs, dtype, config, fp8_recipe, skip_wgrad):
 te_inp_hidden_states = torch.randn(
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_attn_mask = (
 torch.rand(
 (
 1,
 1,
 config.seq_len,
 config.seq_len,
 )
 )
 .cuda()
 .bool()
 )

 if skip_wgrad:
 _disable_wgrads(block)

 use_fp8 = fp8_recipe is not None
 with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
 te_out = block(te_inp_hidden_states)
 loss = te_out.sum()
 loss.backward()
 torch.cuda.synchronize()


def _test_sanity_e2e_bert(block, bs, dtype, config, fp8_recipe, skip_wgrad):
 te_inp_hidden_states = torch.randn(
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()

 if dtype == torch.float32:
 mask_shape = torch.Size([1, 1, config.seq_len, config.seq_len])
 else:
 mask_shape = torch.Size([config.seq_len, bs])

 te_inp_attn_mask = torch.rand(mask_shape).cuda().bool()

Comment on lines

+228
 to 
+233

@timmoon10timmoon10Jun 21, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Why does the FP32 test require a 4D mask while other dtypes use a 2D mask?
@ksivamanksivamanJun 21, 2023

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

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

The FP32 doesn't take the FA path and so uses the PyTorch torch softmax path for padding mask which expects this
@timmoon10timmoon10Jun 22, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Ah, so FA vs PyT is no longer an implementation detail since it expects a different mask format. This makes me think we need an option to explicitly enable or disable FA, and to error out instead of falling back to PyT.
@timmoon10timmoon10Jun 22, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Alternatively, we should expect the mask to be in the format for PyT, and then convert it to the FA format internally so it isn't visible to the user.
@ptrendxptrendxJun 23, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Agree with Tim. OR, accept both in both cases and generate the right one for the current implementation.
To be honest though, the current API for mask type is just bad. For example, what if you want padding AND causal? Then your only option is set the padding mask type and create the causal mask yourself (which would mean you can't do FA in the current implementation which is bad). There is also nothing stopping you really from having arbitrary padding mask (as in, with random elements zeroed out), which would break the assumptions you have in this PR. I think what we could do is introduce new names for the padding type (let's say "pad", "arbitrary" and "no_mask"), deprecate the old names and add causal as separate switch. Then we could say that pad only accepts the "mask" that is actually just list of sequence lengths, arbitrary is whatever you want (and will not go through FA) and causal could be switched irrespective of the mask type.

 if skip_wgrad:
 _disable_wgrads(block)
Expand All@@ -260,26 +247,24 @@ def _test_sanity_e2e_T5(block, bs, dtype, config, fp8_recipe, skip_wgrad):
 te_inp_hidden_states = torch.randn(
 config.seq_len, bs, config.hidden_size, dtype=dtype, requires_grad=True
 ).cuda()
 te_inp_attn_mask = (
 torch.rand(
 (
 1,
 1,
 config.seq_len,
 config.seq_len,
 )
 )
 .cuda()
 .bool()
 )

 if dtype == torch.float32:
 mask_shape = torch.Size([1, 1, config.seq_len, config.seq_len])
 else:
 mask_shape = torch.Size([config.seq_len, bs])

 enc_dec_attn_mask = torch.rand(mask_shape).cuda().bool()

 if skip_wgrad:
 _disable_wgrads(block)

 use_fp8 = fp8_recipe is not None
 with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
 te_out = block(
 te_inp_hidden_states, te_inp_attn_mask, encoder_output=te_inp_hidden_states
 te_inp_hidden_states,
 None,
 encoder_output=te_inp_hidden_states,
 enc_dec_attn_mask=enc_dec_attn_mask,
 )
 loss = te_out.sum()
 loss.backward()
Expand DownExpand Up@@ -468,12 +453,14 @@ def test_sanity_bert(dtype, bs, fp8_recipe, model, skip_wgrad, zero_centered_gam
 apply_residual_connection_post_layernorm=True,
 output_layernorm=True,
 zero_centered_gamma=zero_centered_gamma,
 self_attn_mask_type="padding",

 )
 .to(dtype=dtype)
 .cuda()
 )

_test_sanity_e2e(block, bs, dtype, config, fp8_recipe, skip_wgrad)
_test_sanity_e2e_bert(block, bs, dtype, config, fp8_recipe, skip_wgrad)


@pytest.mark.parametrize("dtype", param_types)
Expand Down

61 changes: 41 additions & 20 deletions
transformer_engine/pytorch/attention.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -44,6 +44,7 @@
 checkpoint,
)
from transformer_engine.pytorch.export import is_in_onnx_export_mode
from transformer_engine.pytorch.jit import jit_fuser

_flash_attn_version = packaging.version.Version(version("flash-attn"))
_flash_attn_version_required = packaging.version.Version("1.0.6")
Expand All@@ -52,6 +53,19 @@
__all__ = ["DotProductAttention"]


@jit_fuser
def get_cu_seqlens(padding_mask: torch.Tensor) -> torch.Tensor:
 """
 Given a padding mask of shape [seq_len, batch_size], returns an int32
@cyanguwacyanguwaJun 22, 2023

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 guess by seq_len you mean max_seqlen? Also, should the mask tensor be required to be a CUDA tensor?
@ksivamanksivamanJun 22, 2023

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

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

Yeah it is the max_seqlen.
@ksivamanksivamanJun 22, 2023

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

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

Yeah it should be cuda tensor since the inp is cuda as well.
 tensor of shape [batch_size + 1,] containing the cumulative sequence
 lengths of every sample in the batch.
 """
 reduced_mask = padding_mask.sum(dim=0)
 cu_seqlens = reduced_mask.cumsum(dim=0).to(torch.int32)
 zero = torch.zeros(1, dtype=torch.int32, device="cuda")
 return torch.cat((zero, cu_seqlens))


def _rotate_half(x: torch.Tensor) -> torch.Tensor:
 """
 change sign so the last dimension becomes [-odd, +even]
Expand DownExpand Up@@ -345,7 +359,7 @@ def __init__(
 _flash_attn_version >= _flash_attn_version_required
 ), f"FlashAttention minimum version {_flash_attn_version_required} is required."

 self.attn_causal_mask = attn_mask_type == "causal"
 self.attn_mask_type = attn_mask_type
 self.norm_factor = norm_factor
 self.attention_dropout_ctx = attention_dropout_ctx
 self.attention_dropout = attention_dropout
Expand All@@ -356,6 +370,7 @@ def forward(
 query_layer: torch.Tensor,
 key_layer: torch.Tensor,
 value_layer: torch.Tensor,
 attention_mask: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
 """flash-attn fprop"""

Expand All@@ -369,7 +384,6 @@ def forward(
 ), 'FlashAttention currently only supports CUDA tensors.'

 # For now just 128, will make it more general in the future

 if (query_layer.shape[-1] == 128 and
 query_layer.shape[0] * query_layer.shape[1] >= 512 and
 _check_if_interleaved_qkv(query_layer, key_layer, value_layer)):
Expand All@@ -389,18 +403,30 @@ def forward(
 ]

 max_seqlen = seqlen
 cu_seqlens = torch.arange(
 0,
 (batch_size + 1) * seqlen,
 step=seqlen,
 dtype=torch.int32,
 device=query_layer.device)
 if self.attn_mask_type == "padding":
 assert (
 attention_mask is not None
 ), "Boolean attention mask must be provided for padding."
@cyanguwacyanguwaJun 22, 2023

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.

So we're settled on requiring a mask tensor than a cu_seqlen tensor from users?
@ksivamanksivamanJun 22, 2023•

edited

Loading

Copy link


Copy Markdown

MemberAuthor
Choose a reason for hiding this comment

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

Yeah I think so, just to be consistent

 mask_shape = torch.Size([seqlen, batch_size])
 assert (
 attention_mask.shape == mask_shape
 ), f"Expected shape {mask_shape} for attenion mask but found {attention_mask.shape}."

 cu_seqlens = get_cu_seqlens(attention_mask)
 else:
 cu_seqlens = torch.arange(
 0,
 (batch_size + 1) * seqlen,
 step=seqlen,
 dtype=torch.int32,
 device=query_layer.device)

 with self.attention_dropout_ctx():
 output = flash_attn_unpadded_func(
 query_layer, key_layer, value_layer, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen,
 self.attention_dropout if self.training else 0.0,
 softmax_scale=1.0/self.norm_factor, causal=self.attn_causal_mask,
 softmax_scale=1.0/self.norm_factor, causal=self.attn_mask_type=="causal",
 deterministic=self.deterministic,
 )

Expand DownExpand Up@@ -696,7 +722,7 @@ class DotProductAttention(torch.nn.Module):
 .. note::

 Argument :attr:`attention_mask` will be ignored in the `forward` call when
 :attr:`attn_mask_type` is set to `"causal"`.
 :attr:`attn_mask_type` is set to `"causal"` or `"no_mask"`.

 .. warning::

Expand All@@ -714,7 +740,7 @@ class DotProductAttention(torch.nn.Module):
 number of key-value channels.
 attention_dropout: float, default = 0.0
 dropout probability for the dropout op during multi-head attention.
 attn_mask_type: {'causal', 'padding'}, default = `causal`
 attn_mask_type: {'causal', 'padding', 'no_mask'}, default = `causal`
 type of attention mask passed into softmax operation.
 layer_number: int, default = `None`
 layer number of the current `DotProductAttention` when multiple such modules
Expand DownExpand Up@@ -829,7 +855,7 @@ def forward(
 .. note::

 Argument :attr:`attention_mask` will be ignored when :attr:`attn_mask_type`
 is set to `"causal"`.
 is set to `"causal"` or `"no_mask"`.

 .. note::

Expand DownExpand Up@@ -884,7 +910,6 @@ def forward(
 use_flash_attention = False

 if self.attn_mask_type == "padding" and attention_mask is not None:
 use_flash_attention = False
 use_fused_attention = False

 if is_in_onnx_export_mode():
Expand All@@ -911,8 +936,9 @@ def forward(
 return self._checkpointed_attention_forward(self.flash_attention,
 query_layer,
 key_layer,
 value_layer)
 return self.flash_attention(query_layer, key_layer, value_layer)
 value_layer,
 attention_mask)
 return self.flash_attention(query_layer, key_layer, value_layer, attention_mask)

 if use_fused_attention:
 if checkpoint_core_attention:
Expand DownExpand Up@@ -1139,11 +1165,6 @@ def forward(
 """MultiHeadAttention FWD"""
 # hidden_states: [sq, b, h]

 if self.attn_mask_type != "causal" and attention_mask is not None:
 assert (
 attention_mask.dtype == torch.bool
 ), "Attention mask must be a boolean tensor"

 # =================================================
 # Pre-allocate memory for key-values for inference.
 # =================================================
Expand Down
Loading