Fix flash attention#84

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Merged
ksivaman merged 5 commits into from
Mar 7, 2023
Merged

Fix flash attention #84

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
f3184f3
ignore self attention mask for causal type
ksivaman Mar 2, 2023
20317c2
further relax checks to run FA, update docs
ksivaman Mar 2, 2023
b0ef5ab
fix pytorch softmax path
ksivaman Mar 2, 2023
23c5616
fixes
ksivaman Mar 6, 2023
f696393
minimum ampere requirement for fa
ksivaman Mar 7, 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
1 change: 1 addition & 0 deletions tests/pytorch/test_onnx_export.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -792,6 +792,7 @@ def test_export_core_attention(

if attn_mask_type is None:
attn_mask_type = 'causal'
inp = (query_layer, key_layer, value_layer)
model = te.transformer.DotProductAttention(
num_attention_heads=num_attention_heads,
kv_channels=kv_channels,
Expand Down
13 changes: 13 additions & 0 deletions transformer_engine/pytorch/softmax.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -16,6 +16,15 @@
THREADS_PER_BLOCK = 128


_default_causal_mask = {}

def _get_default_causal_mask(sq: int) -> torch.Tensor:
"""Return the causal upper triangular mask for softmax input"""
if sq not in _default_causal_mask:
_default_causal_mask[sq] = torch.triu(torch.ones(sq, sq, device="cuda"), diagonal=1).bool()
return _default_causal_mask[sq]


class ScaledUpperTriangMaskedSoftmax(torch.autograd.Function):
"""
Fused operation which performs following three operations in sequence
Expand DownExpand Up@@ -274,6 +283,10 @@ def forward_torch_softmax(

if self.scale is not None:
inp = inp * self.scale

if self.attn_mask_type == "causal":
mask = _get_default_causal_mask(inp.size()[2])

mask_output = self.mask_func(inp, mask) if mask is not None else inp
probs = torch.nn.Softmax(dim=-1)(mask_output)

Expand Down
46 changes: 28 additions & 18 deletions transformer_engine/pytorch/transformer.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -27,6 +27,7 @@
split_tensor_along_dim,
cast_if_needed,
get_default_init_method,
get_device_compute_capability,
)
from transformer_engine.pytorch.constants import (
AttnMaskTypes,
Expand DownExpand Up@@ -220,16 +221,14 @@ def __init__(
assert (
attn_mask_type == "causal"
), 'FlashAttention currently only supports causal attention mask.'
assert (
attention_softmax_in_fp32
), 'FlashAttention currently only supports softmax compute in fp32.'

self.attn_causal_mask = attn_mask_type == "causal"
self.norm_factor = norm_factor
self.attention_dropout_ctx = attention_dropout_ctx
self.attention_dropout = attention_dropout
self.layer_number = layer_number
self.apply_query_key_layer_scaling = apply_query_key_layer_scaling
self.attention_softmax_in_fp32 = attention_softmax_in_fp32

def forward(
self,
Expand DownExpand Up@@ -287,6 +286,11 @@ class DotProductAttention(torch.nn.Module):
representation subspaces as described in the paper:
`Attention Is All You Need <https://arxiv.org/abs/1706.03762>`_.

.. note::

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

.. warning::

For the default attention mechanism, this module executes a non-deterministic version of
Expand All@@ -303,15 +307,6 @@ 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.
layer_number: int, default = `None`
layer number of the current `DotProductAttention` when multiple such modules
are concatenated, for instance in consecutive transformer blocks.
apply_query_key_layer_scaling: bool, default = `False`
apply query-key layer scaling during BMM1
by a factor of `layer_number`
attention_softmax_in_fp32: bool, default = `True`
if set to `False`, softmax is executed in
the dtype of activation tensors.
attn_mask_type: {'causal', 'padding'}, default = `causal`
type of attention mask passed into softmax operation.

Expand DownExpand Up@@ -371,9 +366,8 @@ def __init__(

self.use_flash_attention = (
int(os.getenv("NVTE_FLASH_ATTN", "1"))
and attention_softmax_in_fp32
and attn_mask_type == "causal"
and not apply_query_key_layer_scaling
and get_device_compute_capability() >= 8.0
)

attn_kwargs = {
Expand DownExpand Up@@ -422,6 +416,11 @@ def forward(
"""
Dot Product Attention Layer.

.. note::

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

.. note::

Input tensors :attr:`query_layer`, :attr:`key_layer`, and :attr:`value_layer`
Expand All@@ -448,8 +447,7 @@ def forward(
"""

use_flash_attention = self.use_flash_attention
if (attention_mask is not None
or query_layer.dtype not in [torch.bfloat16, torch.float16]
if (query_layer.dtype not in [torch.bfloat16, torch.float16]
or key_layer.dtype not in [torch.bfloat16, torch.float16]
or value_layer.dtype not in [torch.bfloat16, torch.float16]
):
Expand DownExpand Up@@ -515,6 +513,7 @@ def __init__(
self.return_layernorm_output = return_layernorm_output
self.params_dtype = params_dtype
self.init_method = init_method
self.attn_mask_type = attn_mask_type

if not fuse_qkv_params:
qkv_weight_interleaved = False
Expand DownExpand Up@@ -658,7 +657,7 @@ def forward(
"""MultiHeadAttention FWD"""
# hidden_states: [sq, b, h]

if attention_mask is not None:
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"
Expand DownExpand Up@@ -836,6 +835,11 @@ class TransformerLayer(torch.nn.Module):
TransformerLayer is made up of an attention block and a feedforward network (MLP).
This standard layer is based on the paper "Attention Is All You Need".

.. note::

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

Parameters
----------
hidden_size : int
Expand DownExpand Up@@ -983,6 +987,7 @@ def __init__(
self.apply_residual_connection_post_layernorm = (
apply_residual_connection_post_layernorm
)
self.self_attn_mask_type = self_attn_mask_type
assert (
self_attn_mask_type in AttnMaskTypes
), f"self_attn_mask_type {self_attn_mask_type} not supported"
Expand DownExpand Up@@ -1129,6 +1134,11 @@ def forward(
"""
Transformer Layer: attention block and a feedforward network (MLP)

.. note::

Argument :attr:`attention_mask` will be ignored when :attr:`self_attn_mask_type`
is set to `"causal"`.

Parameters
----------
hidden_states : torch.Tensor
Expand DownExpand Up@@ -1163,7 +1173,7 @@ def forward(

hidden_states = hidden_states.contiguous()

if attention_mask is not None:
if self.self_attn_mask_type != "causal" and attention_mask is not None:
assert (
attention_mask.dtype == torch.bool
), "Attention mask must be a boolean tensor"
Expand Down
7 changes: 7 additions & 0 deletions transformer_engine/pytorch/utils.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -8,6 +8,13 @@
import torch


def get_device_compute_capability() -> float:
"""Returns the cuda compute capability of current GPU"""
major = torch.cuda.get_device_properties(torch.cuda.current_device()).major
minor = torch.cuda.get_device_properties(torch.cuda.current_device()).minor
return major + minor / 10


def attention_mask_func(
attention_scores: torch.Tensor, attention_mask: torch.Tensor
) -> torch.Tensor:
Expand Down
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Add copy buttons to all 
 blocks\n(function() {\n function addCopyButtons() {\n document.querySelectorAll('pre code').forEach(function(codeBlock) {\n if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;\n codeBlock.parentElement.setAttribute('data-copy-added', 'true');\n \n var btn = document.createElement('button');\n btn.textContent = 'Copy';\n 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;';\n btn.onmouseover = function() { this.style.opacity = '1'; };\n btn.onmouseout = function() { this.style.opacity = '0.7'; };\n btn.onclick = function() {\n navigator.clipboard.writeText(codeBlock.textContent).then(function() {\n btn.textContent = 'Copied!';\n setTimeout(function() { btn.textContent = 'Copy'; }, 1500);\n });\n };\n codeBlock.parentElement.style.position = 'relative';\n codeBlock.parentElement.appendChild(btn);\n });\n }\n \n addCopyButtons();\n \n // Re-run on dynamic content\n var observer = new MutationObserver(addCopyButtons);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Add Copy Buttons to Code Blocks");
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '

Skip to content
Fix flash attention#84

New issue

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

 Sign up for GitHub

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

Already on GitHub?
Sign in
to your account

 Merged

ksivaman
merged 5 commits into

from

Mar 7, 2023
 Merged


Fix flash attention
#84
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


5 commits


Select commit
Hold shift + click to select a range
f3184f3

ignore self attention mask for causal type


ksivaman Mar 2, 2023
20317c2

further relax checks to run FA, update docs


ksivaman Mar 2, 2023
b0ef5ab

fix pytorch softmax path


ksivaman Mar 2, 2023
23c5616

fixes


ksivaman Mar 6, 2023
f696393

minimum ampere requirement for fa


ksivaman Mar 7, 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


1 change: 1 addition & 0 deletions
tests/pytorch/test_onnx_export.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -792,6 +792,7 @@ def test_export_core_attention(

 if attn_mask_type is None:
 attn_mask_type = 'causal'
 inp = (query_layer, key_layer, value_layer)
 model = te.transformer.DotProductAttention(
 num_attention_heads=num_attention_heads,
 kv_channels=kv_channels,
Expand Down

13 changes: 13 additions & 0 deletions
transformer_engine/pytorch/softmax.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,6 +16,15 @@
THREADS_PER_BLOCK = 128


_default_causal_mask = {}

def _get_default_causal_mask(sq: int) -> torch.Tensor:
 """Return the causal upper triangular mask for softmax input"""
 if sq not in _default_causal_mask:
 _default_causal_mask[sq] = torch.triu(torch.ones(sq, sq, device="cuda"), diagonal=1).bool()
 return _default_causal_mask[sq]


class ScaledUpperTriangMaskedSoftmax(torch.autograd.Function):
 """
 Fused operation which performs following three operations in sequence
Expand DownExpand Up@@ -274,6 +283,10 @@ def forward_torch_softmax(

 if self.scale is not None:
 inp = inp * self.scale

 if self.attn_mask_type == "causal":
 mask = _get_default_causal_mask(inp.size()[2])

 mask_output = self.mask_func(inp, mask) if mask is not None else inp
 probs = torch.nn.Softmax(dim=-1)(mask_output)

Expand Down

46 changes: 28 additions & 18 deletions
transformer_engine/pytorch/transformer.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -27,6 +27,7 @@
 split_tensor_along_dim,
 cast_if_needed,
 get_default_init_method,
 get_device_compute_capability,
)
from transformer_engine.pytorch.constants import (
 AttnMaskTypes,
Expand DownExpand Up@@ -220,16 +221,14 @@ def __init__(
 assert (
 attn_mask_type == "causal"
 ), 'FlashAttention currently only supports causal attention mask.'
 assert (
 attention_softmax_in_fp32
 ), 'FlashAttention currently only supports softmax compute in fp32.'

 self.attn_causal_mask = attn_mask_type == "causal"
 self.norm_factor = norm_factor
 self.attention_dropout_ctx = attention_dropout_ctx
 self.attention_dropout = attention_dropout
 self.layer_number = layer_number
 self.apply_query_key_layer_scaling = apply_query_key_layer_scaling
 self.attention_softmax_in_fp32 = attention_softmax_in_fp32

 def forward(
 self,
Expand DownExpand Up@@ -287,6 +286,11 @@ class DotProductAttention(torch.nn.Module):
 representation subspaces as described in the paper:
 `Attention Is All You Need <https://arxiv.org/abs/1706.03762>`_.

 .. note::

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

 .. warning::

 For the default attention mechanism, this module executes a non-deterministic version of
Expand All@@ -303,15 +307,6 @@ 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.
 layer_number: int, default = `None`
 layer number of the current `DotProductAttention` when multiple such modules
 are concatenated, for instance in consecutive transformer blocks.
 apply_query_key_layer_scaling: bool, default = `False`
 apply query-key layer scaling during BMM1
 by a factor of `layer_number`
 attention_softmax_in_fp32: bool, default = `True`
 if set to `False`, softmax is executed in
 the dtype of activation tensors.
 attn_mask_type: {'causal', 'padding'}, default = `causal`
 type of attention mask passed into softmax operation.

Expand DownExpand Up@@ -371,9 +366,8 @@ def __init__(

 self.use_flash_attention = (
 int(os.getenv("NVTE_FLASH_ATTN", "1"))
 and attention_softmax_in_fp32
 and attn_mask_type == "causal"
 and not apply_query_key_layer_scaling
 and get_device_compute_capability() >= 8.0
 )

 attn_kwargs = {
Expand DownExpand Up@@ -422,6 +416,11 @@ def forward(
 """
 Dot Product Attention Layer.

 .. note::

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

 .. note::

 Input tensors :attr:`query_layer`, :attr:`key_layer`, and :attr:`value_layer`
Expand All@@ -448,8 +447,7 @@ def forward(
 """

 use_flash_attention = self.use_flash_attention
 if (attention_mask is not None
 or query_layer.dtype not in [torch.bfloat16, torch.float16]
 if (query_layer.dtype not in [torch.bfloat16, torch.float16]
 or key_layer.dtype not in [torch.bfloat16, torch.float16]
 or value_layer.dtype not in [torch.bfloat16, torch.float16]
 ):
Expand DownExpand Up@@ -515,6 +513,7 @@ def __init__(
 self.return_layernorm_output = return_layernorm_output
 self.params_dtype = params_dtype
 self.init_method = init_method
 self.attn_mask_type = attn_mask_type

 if not fuse_qkv_params:
 qkv_weight_interleaved = False
Expand DownExpand Up@@ -658,7 +657,7 @@ def forward(
 """MultiHeadAttention FWD"""
 # hidden_states: [sq, b, h]

 if attention_mask is not None:
 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"
Expand DownExpand Up@@ -836,6 +835,11 @@ class TransformerLayer(torch.nn.Module):
 TransformerLayer is made up of an attention block and a feedforward network (MLP).
 This standard layer is based on the paper "Attention Is All You Need".

 .. note::

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

 Parameters
 ----------
 hidden_size : int
Expand DownExpand Up@@ -983,6 +987,7 @@ def __init__(
 self.apply_residual_connection_post_layernorm = (
 apply_residual_connection_post_layernorm
 )
 self.self_attn_mask_type = self_attn_mask_type
 assert (
 self_attn_mask_type in AttnMaskTypes
 ), f"self_attn_mask_type {self_attn_mask_type} not supported"
Expand DownExpand Up@@ -1129,6 +1134,11 @@ def forward(
 """
 Transformer Layer: attention block and a feedforward network (MLP)

 .. note::

 Argument :attr:`attention_mask` will be ignored when :attr:`self_attn_mask_type`
 is set to `"causal"`.

 Parameters
 ----------
 hidden_states : torch.Tensor
Expand DownExpand Up@@ -1163,7 +1173,7 @@ def forward(

 hidden_states = hidden_states.contiguous()

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

7 changes: 7 additions & 0 deletions
transformer_engine/pytorch/utils.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -8,6 +8,13 @@
import torch


def get_device_compute_capability() -> float:
 """Returns the cuda compute capability of current GPU"""
 major = torch.cuda.get_device_properties(torch.cuda.current_device()).major
 minor = torch.cuda.get_device_properties(torch.cuda.current_device()).minor
 return major + minor / 10


def attention_mask_func(
 attention_scores: torch.Tensor, attention_mask: torch.Tensor
) -> torch.Tensor:
Expand Down

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

Skip to content
Fix flash attention#84

New issue

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

 Sign up for GitHub

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

Already on GitHub?
Sign in
to your account

 Merged

ksivaman
merged 5 commits into

from

Mar 7, 2023
 Merged


Fix flash attention
#84
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


5 commits


Select commit
Hold shift + click to select a range
f3184f3

ignore self attention mask for causal type


ksivaman Mar 2, 2023
20317c2

further relax checks to run FA, update docs


ksivaman Mar 2, 2023
b0ef5ab

fix pytorch softmax path


ksivaman Mar 2, 2023
23c5616

fixes


ksivaman Mar 6, 2023
f696393

minimum ampere requirement for fa


ksivaman Mar 7, 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


1 change: 1 addition & 0 deletions
tests/pytorch/test_onnx_export.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -792,6 +792,7 @@ def test_export_core_attention(

 if attn_mask_type is None:
 attn_mask_type = 'causal'
 inp = (query_layer, key_layer, value_layer)
 model = te.transformer.DotProductAttention(
 num_attention_heads=num_attention_heads,
 kv_channels=kv_channels,
Expand Down

13 changes: 13 additions & 0 deletions
transformer_engine/pytorch/softmax.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,6 +16,15 @@
THREADS_PER_BLOCK = 128


_default_causal_mask = {}

def _get_default_causal_mask(sq: int) -> torch.Tensor:
 """Return the causal upper triangular mask for softmax input"""
 if sq not in _default_causal_mask:
 _default_causal_mask[sq] = torch.triu(torch.ones(sq, sq, device="cuda"), diagonal=1).bool()
 return _default_causal_mask[sq]


class ScaledUpperTriangMaskedSoftmax(torch.autograd.Function):
 """
 Fused operation which performs following three operations in sequence
Expand DownExpand Up@@ -274,6 +283,10 @@ def forward_torch_softmax(

 if self.scale is not None:
 inp = inp * self.scale

 if self.attn_mask_type == "causal":
 mask = _get_default_causal_mask(inp.size()[2])

 mask_output = self.mask_func(inp, mask) if mask is not None else inp
 probs = torch.nn.Softmax(dim=-1)(mask_output)

Expand Down

46 changes: 28 additions & 18 deletions
transformer_engine/pytorch/transformer.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -27,6 +27,7 @@
 split_tensor_along_dim,
 cast_if_needed,
 get_default_init_method,
 get_device_compute_capability,
)
from transformer_engine.pytorch.constants import (
 AttnMaskTypes,
Expand DownExpand Up@@ -220,16 +221,14 @@ def __init__(
 assert (
 attn_mask_type == "causal"
 ), 'FlashAttention currently only supports causal attention mask.'
 assert (
 attention_softmax_in_fp32
 ), 'FlashAttention currently only supports softmax compute in fp32.'

 self.attn_causal_mask = attn_mask_type == "causal"
 self.norm_factor = norm_factor
 self.attention_dropout_ctx = attention_dropout_ctx
 self.attention_dropout = attention_dropout
 self.layer_number = layer_number
 self.apply_query_key_layer_scaling = apply_query_key_layer_scaling
 self.attention_softmax_in_fp32 = attention_softmax_in_fp32

 def forward(
 self,
Expand DownExpand Up@@ -287,6 +286,11 @@ class DotProductAttention(torch.nn.Module):
 representation subspaces as described in the paper:
 `Attention Is All You Need <https://arxiv.org/abs/1706.03762>`_.

 .. note::

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

 .. warning::

 For the default attention mechanism, this module executes a non-deterministic version of
Expand All@@ -303,15 +307,6 @@ 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.
 layer_number: int, default = `None`
 layer number of the current `DotProductAttention` when multiple such modules
 are concatenated, for instance in consecutive transformer blocks.
 apply_query_key_layer_scaling: bool, default = `False`
 apply query-key layer scaling during BMM1
 by a factor of `layer_number`
 attention_softmax_in_fp32: bool, default = `True`
 if set to `False`, softmax is executed in
 the dtype of activation tensors.
 attn_mask_type: {'causal', 'padding'}, default = `causal`
 type of attention mask passed into softmax operation.

Expand DownExpand Up@@ -371,9 +366,8 @@ def __init__(

 self.use_flash_attention = (
 int(os.getenv("NVTE_FLASH_ATTN", "1"))
 and attention_softmax_in_fp32
 and attn_mask_type == "causal"
 and not apply_query_key_layer_scaling
 and get_device_compute_capability() >= 8.0
 )

 attn_kwargs = {
Expand DownExpand Up@@ -422,6 +416,11 @@ def forward(
 """
 Dot Product Attention Layer.

 .. note::

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

 .. note::

 Input tensors :attr:`query_layer`, :attr:`key_layer`, and :attr:`value_layer`
Expand All@@ -448,8 +447,7 @@ def forward(
 """

 use_flash_attention = self.use_flash_attention
 if (attention_mask is not None
 or query_layer.dtype not in [torch.bfloat16, torch.float16]
 if (query_layer.dtype not in [torch.bfloat16, torch.float16]
 or key_layer.dtype not in [torch.bfloat16, torch.float16]
 or value_layer.dtype not in [torch.bfloat16, torch.float16]
 ):
Expand DownExpand Up@@ -515,6 +513,7 @@ def __init__(
 self.return_layernorm_output = return_layernorm_output
 self.params_dtype = params_dtype
 self.init_method = init_method
 self.attn_mask_type = attn_mask_type

 if not fuse_qkv_params:
 qkv_weight_interleaved = False
Expand DownExpand Up@@ -658,7 +657,7 @@ def forward(
 """MultiHeadAttention FWD"""
 # hidden_states: [sq, b, h]

 if attention_mask is not None:
 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"
Expand DownExpand Up@@ -836,6 +835,11 @@ class TransformerLayer(torch.nn.Module):
 TransformerLayer is made up of an attention block and a feedforward network (MLP).
 This standard layer is based on the paper "Attention Is All You Need".

 .. note::

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

 Parameters
 ----------
 hidden_size : int
Expand DownExpand Up@@ -983,6 +987,7 @@ def __init__(
 self.apply_residual_connection_post_layernorm = (
 apply_residual_connection_post_layernorm
 )
 self.self_attn_mask_type = self_attn_mask_type
 assert (
 self_attn_mask_type in AttnMaskTypes
 ), f"self_attn_mask_type {self_attn_mask_type} not supported"
Expand DownExpand Up@@ -1129,6 +1134,11 @@ def forward(
 """
 Transformer Layer: attention block and a feedforward network (MLP)

 .. note::

 Argument :attr:`attention_mask` will be ignored when :attr:`self_attn_mask_type`
 is set to `"causal"`.

 Parameters
 ----------
 hidden_states : torch.Tensor
Expand DownExpand Up@@ -1163,7 +1173,7 @@ def forward(

 hidden_states = hidden_states.contiguous()

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

7 changes: 7 additions & 0 deletions
transformer_engine/pytorch/utils.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -8,6 +8,13 @@
import torch


def get_device_compute_capability() -> float:
 """Returns the cuda compute capability of current GPU"""
 major = torch.cuda.get_device_properties(torch.cuda.current_device()).major
 minor = torch.cuda.get_device_properties(torch.cuda.current_device()).minor
 return major + minor / 10


def attention_mask_func(
 attention_scores: torch.Tensor, attention_mask: torch.Tensor
) -> torch.Tensor:
Expand Down

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

Skip to content
Fix flash attention#84

New issue

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

 Sign up for GitHub

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

Already on GitHub?
Sign in
to your account

 Merged

ksivaman
merged 5 commits into

from

Mar 7, 2023
 Merged


Fix flash attention
#84
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


5 commits


Select commit
Hold shift + click to select a range
f3184f3

ignore self attention mask for causal type


ksivaman Mar 2, 2023
20317c2

further relax checks to run FA, update docs


ksivaman Mar 2, 2023
b0ef5ab

fix pytorch softmax path


ksivaman Mar 2, 2023
23c5616

fixes


ksivaman Mar 6, 2023
f696393

minimum ampere requirement for fa


ksivaman Mar 7, 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


1 change: 1 addition & 0 deletions
tests/pytorch/test_onnx_export.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -792,6 +792,7 @@ def test_export_core_attention(

 if attn_mask_type is None:
 attn_mask_type = 'causal'
 inp = (query_layer, key_layer, value_layer)
 model = te.transformer.DotProductAttention(
 num_attention_heads=num_attention_heads,
 kv_channels=kv_channels,
Expand Down

13 changes: 13 additions & 0 deletions
transformer_engine/pytorch/softmax.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,6 +16,15 @@
THREADS_PER_BLOCK = 128


_default_causal_mask = {}

def _get_default_causal_mask(sq: int) -> torch.Tensor:
 """Return the causal upper triangular mask for softmax input"""
 if sq not in _default_causal_mask:
 _default_causal_mask[sq] = torch.triu(torch.ones(sq, sq, device="cuda"), diagonal=1).bool()
 return _default_causal_mask[sq]


class ScaledUpperTriangMaskedSoftmax(torch.autograd.Function):
 """
 Fused operation which performs following three operations in sequence
Expand DownExpand Up@@ -274,6 +283,10 @@ def forward_torch_softmax(

 if self.scale is not None:
 inp = inp * self.scale

 if self.attn_mask_type == "causal":
 mask = _get_default_causal_mask(inp.size()[2])

 mask_output = self.mask_func(inp, mask) if mask is not None else inp
 probs = torch.nn.Softmax(dim=-1)(mask_output)

Expand Down

46 changes: 28 additions & 18 deletions
transformer_engine/pytorch/transformer.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -27,6 +27,7 @@
 split_tensor_along_dim,
 cast_if_needed,
 get_default_init_method,
 get_device_compute_capability,
)
from transformer_engine.pytorch.constants import (
 AttnMaskTypes,
Expand DownExpand Up@@ -220,16 +221,14 @@ def __init__(
 assert (
 attn_mask_type == "causal"
 ), 'FlashAttention currently only supports causal attention mask.'
 assert (
 attention_softmax_in_fp32
 ), 'FlashAttention currently only supports softmax compute in fp32.'

 self.attn_causal_mask = attn_mask_type == "causal"
 self.norm_factor = norm_factor
 self.attention_dropout_ctx = attention_dropout_ctx
 self.attention_dropout = attention_dropout
 self.layer_number = layer_number
 self.apply_query_key_layer_scaling = apply_query_key_layer_scaling
 self.attention_softmax_in_fp32 = attention_softmax_in_fp32

 def forward(
 self,
Expand DownExpand Up@@ -287,6 +286,11 @@ class DotProductAttention(torch.nn.Module):
 representation subspaces as described in the paper:
 `Attention Is All You Need <https://arxiv.org/abs/1706.03762>`_.

 .. note::

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

 .. warning::

 For the default attention mechanism, this module executes a non-deterministic version of
Expand All@@ -303,15 +307,6 @@ 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.
 layer_number: int, default = `None`
 layer number of the current `DotProductAttention` when multiple such modules
 are concatenated, for instance in consecutive transformer blocks.
 apply_query_key_layer_scaling: bool, default = `False`
 apply query-key layer scaling during BMM1
 by a factor of `layer_number`
 attention_softmax_in_fp32: bool, default = `True`
 if set to `False`, softmax is executed in
 the dtype of activation tensors.
 attn_mask_type: {'causal', 'padding'}, default = `causal`
 type of attention mask passed into softmax operation.

Expand DownExpand Up@@ -371,9 +366,8 @@ def __init__(

 self.use_flash_attention = (
 int(os.getenv("NVTE_FLASH_ATTN", "1"))
 and attention_softmax_in_fp32
 and attn_mask_type == "causal"
 and not apply_query_key_layer_scaling
 and get_device_compute_capability() >= 8.0
 )

 attn_kwargs = {
Expand DownExpand Up@@ -422,6 +416,11 @@ def forward(
 """
 Dot Product Attention Layer.

 .. note::

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

 .. note::

 Input tensors :attr:`query_layer`, :attr:`key_layer`, and :attr:`value_layer`
Expand All@@ -448,8 +447,7 @@ def forward(
 """

 use_flash_attention = self.use_flash_attention
 if (attention_mask is not None
 or query_layer.dtype not in [torch.bfloat16, torch.float16]
 if (query_layer.dtype not in [torch.bfloat16, torch.float16]
 or key_layer.dtype not in [torch.bfloat16, torch.float16]
 or value_layer.dtype not in [torch.bfloat16, torch.float16]
 ):
Expand DownExpand Up@@ -515,6 +513,7 @@ def __init__(
 self.return_layernorm_output = return_layernorm_output
 self.params_dtype = params_dtype
 self.init_method = init_method
 self.attn_mask_type = attn_mask_type

 if not fuse_qkv_params:
 qkv_weight_interleaved = False
Expand DownExpand Up@@ -658,7 +657,7 @@ def forward(
 """MultiHeadAttention FWD"""
 # hidden_states: [sq, b, h]

 if attention_mask is not None:
 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"
Expand DownExpand Up@@ -836,6 +835,11 @@ class TransformerLayer(torch.nn.Module):
 TransformerLayer is made up of an attention block and a feedforward network (MLP).
 This standard layer is based on the paper "Attention Is All You Need".

 .. note::

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

 Parameters
 ----------
 hidden_size : int
Expand DownExpand Up@@ -983,6 +987,7 @@ def __init__(
 self.apply_residual_connection_post_layernorm = (
 apply_residual_connection_post_layernorm
 )
 self.self_attn_mask_type = self_attn_mask_type
 assert (
 self_attn_mask_type in AttnMaskTypes
 ), f"self_attn_mask_type {self_attn_mask_type} not supported"
Expand DownExpand Up@@ -1129,6 +1134,11 @@ def forward(
 """
 Transformer Layer: attention block and a feedforward network (MLP)

 .. note::

 Argument :attr:`attention_mask` will be ignored when :attr:`self_attn_mask_type`
 is set to `"causal"`.

 Parameters
 ----------
 hidden_states : torch.Tensor
Expand DownExpand Up@@ -1163,7 +1173,7 @@ def forward(

 hidden_states = hidden_states.contiguous()

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

7 changes: 7 additions & 0 deletions
transformer_engine/pytorch/utils.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -8,6 +8,13 @@
import torch


def get_device_compute_capability() -> float:
 """Returns the cuda compute capability of current GPU"""
 major = torch.cuda.get_device_properties(torch.cuda.current_device()).major
 minor = torch.cuda.get_device_properties(torch.cuda.current_device()).minor
 return major + minor / 10


def attention_mask_func(
 attention_scores: torch.Tensor, attention_mask: torch.Tensor
) -> torch.Tensor:
Expand Down

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

Skip to content
Fix flash attention#84

New issue

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

 Sign up for GitHub

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

Already on GitHub?
Sign in
to your account

 Merged

ksivaman
merged 5 commits into

from

Mar 7, 2023
 Merged


Fix flash attention
#84
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


5 commits


Select commit
Hold shift + click to select a range
f3184f3

ignore self attention mask for causal type


ksivaman Mar 2, 2023
20317c2

further relax checks to run FA, update docs


ksivaman Mar 2, 2023
b0ef5ab

fix pytorch softmax path


ksivaman Mar 2, 2023
23c5616

fixes


ksivaman Mar 6, 2023
f696393

minimum ampere requirement for fa


ksivaman Mar 7, 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


1 change: 1 addition & 0 deletions
tests/pytorch/test_onnx_export.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -792,6 +792,7 @@ def test_export_core_attention(

 if attn_mask_type is None:
 attn_mask_type = 'causal'
 inp = (query_layer, key_layer, value_layer)
 model = te.transformer.DotProductAttention(
 num_attention_heads=num_attention_heads,
 kv_channels=kv_channels,
Expand Down

13 changes: 13 additions & 0 deletions
transformer_engine/pytorch/softmax.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,6 +16,15 @@
THREADS_PER_BLOCK = 128


_default_causal_mask = {}

def _get_default_causal_mask(sq: int) -> torch.Tensor:
 """Return the causal upper triangular mask for softmax input"""
 if sq not in _default_causal_mask:
 _default_causal_mask[sq] = torch.triu(torch.ones(sq, sq, device="cuda"), diagonal=1).bool()
 return _default_causal_mask[sq]


class ScaledUpperTriangMaskedSoftmax(torch.autograd.Function):
 """
 Fused operation which performs following three operations in sequence
Expand DownExpand Up@@ -274,6 +283,10 @@ def forward_torch_softmax(

 if self.scale is not None:
 inp = inp * self.scale

 if self.attn_mask_type == "causal":
 mask = _get_default_causal_mask(inp.size()[2])

 mask_output = self.mask_func(inp, mask) if mask is not None else inp
 probs = torch.nn.Softmax(dim=-1)(mask_output)

Expand Down

46 changes: 28 additions & 18 deletions
transformer_engine/pytorch/transformer.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -27,6 +27,7 @@
 split_tensor_along_dim,
 cast_if_needed,
 get_default_init_method,
 get_device_compute_capability,
)
from transformer_engine.pytorch.constants import (
 AttnMaskTypes,
Expand DownExpand Up@@ -220,16 +221,14 @@ def __init__(
 assert (
 attn_mask_type == "causal"
 ), 'FlashAttention currently only supports causal attention mask.'
 assert (
 attention_softmax_in_fp32
 ), 'FlashAttention currently only supports softmax compute in fp32.'

 self.attn_causal_mask = attn_mask_type == "causal"
 self.norm_factor = norm_factor
 self.attention_dropout_ctx = attention_dropout_ctx
 self.attention_dropout = attention_dropout
 self.layer_number = layer_number
 self.apply_query_key_layer_scaling = apply_query_key_layer_scaling
 self.attention_softmax_in_fp32 = attention_softmax_in_fp32

 def forward(
 self,
Expand DownExpand Up@@ -287,6 +286,11 @@ class DotProductAttention(torch.nn.Module):
 representation subspaces as described in the paper:
 `Attention Is All You Need <https://arxiv.org/abs/1706.03762>`_.

 .. note::

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

 .. warning::

 For the default attention mechanism, this module executes a non-deterministic version of
Expand All@@ -303,15 +307,6 @@ 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.
 layer_number: int, default = `None`
 layer number of the current `DotProductAttention` when multiple such modules
 are concatenated, for instance in consecutive transformer blocks.
 apply_query_key_layer_scaling: bool, default = `False`
 apply query-key layer scaling during BMM1
 by a factor of `layer_number`
 attention_softmax_in_fp32: bool, default = `True`
 if set to `False`, softmax is executed in
 the dtype of activation tensors.
 attn_mask_type: {'causal', 'padding'}, default = `causal`
 type of attention mask passed into softmax operation.

Expand DownExpand Up@@ -371,9 +366,8 @@ def __init__(

 self.use_flash_attention = (
 int(os.getenv("NVTE_FLASH_ATTN", "1"))
 and attention_softmax_in_fp32
 and attn_mask_type == "causal"
 and not apply_query_key_layer_scaling
 and get_device_compute_capability() >= 8.0
 )

 attn_kwargs = {
Expand DownExpand Up@@ -422,6 +416,11 @@ def forward(
 """
 Dot Product Attention Layer.

 .. note::

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

 .. note::

 Input tensors :attr:`query_layer`, :attr:`key_layer`, and :attr:`value_layer`
Expand All@@ -448,8 +447,7 @@ def forward(
 """

 use_flash_attention = self.use_flash_attention
 if (attention_mask is not None
 or query_layer.dtype not in [torch.bfloat16, torch.float16]
 if (query_layer.dtype not in [torch.bfloat16, torch.float16]
 or key_layer.dtype not in [torch.bfloat16, torch.float16]
 or value_layer.dtype not in [torch.bfloat16, torch.float16]
 ):
Expand DownExpand Up@@ -515,6 +513,7 @@ def __init__(
 self.return_layernorm_output = return_layernorm_output
 self.params_dtype = params_dtype
 self.init_method = init_method
 self.attn_mask_type = attn_mask_type

 if not fuse_qkv_params:
 qkv_weight_interleaved = False
Expand DownExpand Up@@ -658,7 +657,7 @@ def forward(
 """MultiHeadAttention FWD"""
 # hidden_states: [sq, b, h]

 if attention_mask is not None:
 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"
Expand DownExpand Up@@ -836,6 +835,11 @@ class TransformerLayer(torch.nn.Module):
 TransformerLayer is made up of an attention block and a feedforward network (MLP).
 This standard layer is based on the paper "Attention Is All You Need".

 .. note::

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

 Parameters
 ----------
 hidden_size : int
Expand DownExpand Up@@ -983,6 +987,7 @@ def __init__(
 self.apply_residual_connection_post_layernorm = (
 apply_residual_connection_post_layernorm
 )
 self.self_attn_mask_type = self_attn_mask_type
 assert (
 self_attn_mask_type in AttnMaskTypes
 ), f"self_attn_mask_type {self_attn_mask_type} not supported"
Expand DownExpand Up@@ -1129,6 +1134,11 @@ def forward(
 """
 Transformer Layer: attention block and a feedforward network (MLP)

 .. note::

 Argument :attr:`attention_mask` will be ignored when :attr:`self_attn_mask_type`
 is set to `"causal"`.

 Parameters
 ----------
 hidden_states : torch.Tensor
Expand DownExpand Up@@ -1163,7 +1173,7 @@ def forward(

 hidden_states = hidden_states.contiguous()

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

7 changes: 7 additions & 0 deletions
transformer_engine/pytorch/utils.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -8,6 +8,13 @@
import torch


def get_device_compute_capability() -> float:
 """Returns the cuda compute capability of current GPU"""
 major = torch.cuda.get_device_properties(torch.cuda.current_device()).major
 minor = torch.cuda.get_device_properties(torch.cuda.current_device()).minor
 return major + minor / 10


def attention_mask_func(
 attention_scores: torch.Tensor, attention_mask: torch.Tensor
) -> torch.Tensor:
Expand Down

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

Skip to content
Fix flash attention#84

New issue

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

 Sign up for GitHub

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

Already on GitHub?
Sign in
to your account

 Merged

ksivaman
merged 5 commits into

from

Mar 7, 2023
 Merged


Fix flash attention
#84
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


5 commits


Select commit
Hold shift + click to select a range
f3184f3

ignore self attention mask for causal type


ksivaman Mar 2, 2023
20317c2

further relax checks to run FA, update docs


ksivaman Mar 2, 2023
b0ef5ab

fix pytorch softmax path


ksivaman Mar 2, 2023
23c5616

fixes


ksivaman Mar 6, 2023
f696393

minimum ampere requirement for fa


ksivaman Mar 7, 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


1 change: 1 addition & 0 deletions
tests/pytorch/test_onnx_export.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -792,6 +792,7 @@ def test_export_core_attention(

 if attn_mask_type is None:
 attn_mask_type = 'causal'
 inp = (query_layer, key_layer, value_layer)
 model = te.transformer.DotProductAttention(
 num_attention_heads=num_attention_heads,
 kv_channels=kv_channels,
Expand Down

13 changes: 13 additions & 0 deletions
transformer_engine/pytorch/softmax.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,6 +16,15 @@
THREADS_PER_BLOCK = 128


_default_causal_mask = {}

def _get_default_causal_mask(sq: int) -> torch.Tensor:
 """Return the causal upper triangular mask for softmax input"""
 if sq not in _default_causal_mask:
 _default_causal_mask[sq] = torch.triu(torch.ones(sq, sq, device="cuda"), diagonal=1).bool()
 return _default_causal_mask[sq]


class ScaledUpperTriangMaskedSoftmax(torch.autograd.Function):
 """
 Fused operation which performs following three operations in sequence
Expand DownExpand Up@@ -274,6 +283,10 @@ def forward_torch_softmax(

 if self.scale is not None:
 inp = inp * self.scale

 if self.attn_mask_type == "causal":
 mask = _get_default_causal_mask(inp.size()[2])

 mask_output = self.mask_func(inp, mask) if mask is not None else inp
 probs = torch.nn.Softmax(dim=-1)(mask_output)

Expand Down

46 changes: 28 additions & 18 deletions
transformer_engine/pytorch/transformer.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -27,6 +27,7 @@
 split_tensor_along_dim,
 cast_if_needed,
 get_default_init_method,
 get_device_compute_capability,
)
from transformer_engine.pytorch.constants import (
 AttnMaskTypes,
Expand DownExpand Up@@ -220,16 +221,14 @@ def __init__(
 assert (
 attn_mask_type == "causal"
 ), 'FlashAttention currently only supports causal attention mask.'
 assert (
 attention_softmax_in_fp32
 ), 'FlashAttention currently only supports softmax compute in fp32.'

 self.attn_causal_mask = attn_mask_type == "causal"
 self.norm_factor = norm_factor
 self.attention_dropout_ctx = attention_dropout_ctx
 self.attention_dropout = attention_dropout
 self.layer_number = layer_number
 self.apply_query_key_layer_scaling = apply_query_key_layer_scaling
 self.attention_softmax_in_fp32 = attention_softmax_in_fp32

 def forward(
 self,
Expand DownExpand Up@@ -287,6 +286,11 @@ class DotProductAttention(torch.nn.Module):
 representation subspaces as described in the paper:
 `Attention Is All You Need <https://arxiv.org/abs/1706.03762>`_.

 .. note::

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

 .. warning::

 For the default attention mechanism, this module executes a non-deterministic version of
Expand All@@ -303,15 +307,6 @@ 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.
 layer_number: int, default = `None`
 layer number of the current `DotProductAttention` when multiple such modules
 are concatenated, for instance in consecutive transformer blocks.
 apply_query_key_layer_scaling: bool, default = `False`
 apply query-key layer scaling during BMM1
 by a factor of `layer_number`
 attention_softmax_in_fp32: bool, default = `True`
 if set to `False`, softmax is executed in
 the dtype of activation tensors.
 attn_mask_type: {'causal', 'padding'}, default = `causal`
 type of attention mask passed into softmax operation.

Expand DownExpand Up@@ -371,9 +366,8 @@ def __init__(

 self.use_flash_attention = (
 int(os.getenv("NVTE_FLASH_ATTN", "1"))
 and attention_softmax_in_fp32
 and attn_mask_type == "causal"
 and not apply_query_key_layer_scaling
 and get_device_compute_capability() >= 8.0
 )

 attn_kwargs = {
Expand DownExpand Up@@ -422,6 +416,11 @@ def forward(
 """
 Dot Product Attention Layer.

 .. note::

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

 .. note::

 Input tensors :attr:`query_layer`, :attr:`key_layer`, and :attr:`value_layer`
Expand All@@ -448,8 +447,7 @@ def forward(
 """

 use_flash_attention = self.use_flash_attention
 if (attention_mask is not None
 or query_layer.dtype not in [torch.bfloat16, torch.float16]
 if (query_layer.dtype not in [torch.bfloat16, torch.float16]
 or key_layer.dtype not in [torch.bfloat16, torch.float16]
 or value_layer.dtype not in [torch.bfloat16, torch.float16]
 ):
Expand DownExpand Up@@ -515,6 +513,7 @@ def __init__(
 self.return_layernorm_output = return_layernorm_output
 self.params_dtype = params_dtype
 self.init_method = init_method
 self.attn_mask_type = attn_mask_type

 if not fuse_qkv_params:
 qkv_weight_interleaved = False
Expand DownExpand Up@@ -658,7 +657,7 @@ def forward(
 """MultiHeadAttention FWD"""
 # hidden_states: [sq, b, h]

 if attention_mask is not None:
 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"
Expand DownExpand Up@@ -836,6 +835,11 @@ class TransformerLayer(torch.nn.Module):
 TransformerLayer is made up of an attention block and a feedforward network (MLP).
 This standard layer is based on the paper "Attention Is All You Need".

 .. note::

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

 Parameters
 ----------
 hidden_size : int
Expand DownExpand Up@@ -983,6 +987,7 @@ def __init__(
 self.apply_residual_connection_post_layernorm = (
 apply_residual_connection_post_layernorm
 )
 self.self_attn_mask_type = self_attn_mask_type
 assert (
 self_attn_mask_type in AttnMaskTypes
 ), f"self_attn_mask_type {self_attn_mask_type} not supported"
Expand DownExpand Up@@ -1129,6 +1134,11 @@ def forward(
 """
 Transformer Layer: attention block and a feedforward network (MLP)

 .. note::

 Argument :attr:`attention_mask` will be ignored when :attr:`self_attn_mask_type`
 is set to `"causal"`.

 Parameters
 ----------
 hidden_states : torch.Tensor
Expand DownExpand Up@@ -1163,7 +1173,7 @@ def forward(

 hidden_states = hidden_states.contiguous()

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

7 changes: 7 additions & 0 deletions
transformer_engine/pytorch/utils.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -8,6 +8,13 @@
import torch


def get_device_compute_capability() -> float:
 """Returns the cuda compute capability of current GPU"""
 major = torch.cuda.get_device_properties(torch.cuda.current_device()).major
 minor = torch.cuda.get_device_properties(torch.cuda.current_device()).minor
 return major + minor / 10


def attention_mask_func(
 attention_scores: torch.Tensor, attention_mask: torch.Tensor
) -> torch.Tensor:
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Remove or un-stick sticky/fixed headers that block content\n(function() {\n function unstick() {\n document.querySelectorAll('header, nav, [role=\"banner\"], .header, .navbar, .sticky, .fixed-top, [style*=\"position: fixed\"], [style*=\"position:sticky\"]').forEach(function(el) {\n if (el.style.position === 'fixed' || el.style.position === 'sticky' || \n getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {\n el.style.position = 'static';\n el.style.top = 'auto';\n el.style.zIndex = 'auto';\n }\n });\n }\n \n unstick();\n \n var observer = new MutationObserver(unstick);\n observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });\n})();", "Kill Sticky Headers");
}
} catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
Fix flash attention#84

New issue

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

 Sign up for GitHub

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

Already on GitHub?
Sign in
to your account

 Merged

ksivaman
merged 5 commits into

from

Mar 7, 2023
 Merged


Fix flash attention
#84
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


5 commits


Select commit
Hold shift + click to select a range
f3184f3

ignore self attention mask for causal type


ksivaman Mar 2, 2023
20317c2

further relax checks to run FA, update docs


ksivaman Mar 2, 2023
b0ef5ab

fix pytorch softmax path


ksivaman Mar 2, 2023
23c5616

fixes


ksivaman Mar 6, 2023
f696393

minimum ampere requirement for fa


ksivaman Mar 7, 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


1 change: 1 addition & 0 deletions
tests/pytorch/test_onnx_export.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -792,6 +792,7 @@ def test_export_core_attention(

 if attn_mask_type is None:
 attn_mask_type = 'causal'
 inp = (query_layer, key_layer, value_layer)
 model = te.transformer.DotProductAttention(
 num_attention_heads=num_attention_heads,
 kv_channels=kv_channels,
Expand Down

13 changes: 13 additions & 0 deletions
transformer_engine/pytorch/softmax.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,6 +16,15 @@
THREADS_PER_BLOCK = 128


_default_causal_mask = {}

def _get_default_causal_mask(sq: int) -> torch.Tensor:
 """Return the causal upper triangular mask for softmax input"""
 if sq not in _default_causal_mask:
 _default_causal_mask[sq] = torch.triu(torch.ones(sq, sq, device="cuda"), diagonal=1).bool()
 return _default_causal_mask[sq]


class ScaledUpperTriangMaskedSoftmax(torch.autograd.Function):
 """
 Fused operation which performs following three operations in sequence
Expand DownExpand Up@@ -274,6 +283,10 @@ def forward_torch_softmax(

 if self.scale is not None:
 inp = inp * self.scale

 if self.attn_mask_type == "causal":
 mask = _get_default_causal_mask(inp.size()[2])

 mask_output = self.mask_func(inp, mask) if mask is not None else inp
 probs = torch.nn.Softmax(dim=-1)(mask_output)

Expand Down

46 changes: 28 additions & 18 deletions
transformer_engine/pytorch/transformer.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -27,6 +27,7 @@
 split_tensor_along_dim,
 cast_if_needed,
 get_default_init_method,
 get_device_compute_capability,
)
from transformer_engine.pytorch.constants import (
 AttnMaskTypes,
Expand DownExpand Up@@ -220,16 +221,14 @@ def __init__(
 assert (
 attn_mask_type == "causal"
 ), 'FlashAttention currently only supports causal attention mask.'
 assert (
 attention_softmax_in_fp32
 ), 'FlashAttention currently only supports softmax compute in fp32.'

 self.attn_causal_mask = attn_mask_type == "causal"
 self.norm_factor = norm_factor
 self.attention_dropout_ctx = attention_dropout_ctx
 self.attention_dropout = attention_dropout
 self.layer_number = layer_number
 self.apply_query_key_layer_scaling = apply_query_key_layer_scaling
 self.attention_softmax_in_fp32 = attention_softmax_in_fp32

 def forward(
 self,
Expand DownExpand Up@@ -287,6 +286,11 @@ class DotProductAttention(torch.nn.Module):
 representation subspaces as described in the paper:
 `Attention Is All You Need <https://arxiv.org/abs/1706.03762>`_.

 .. note::

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

 .. warning::

 For the default attention mechanism, this module executes a non-deterministic version of
Expand All@@ -303,15 +307,6 @@ 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.
 layer_number: int, default = `None`
 layer number of the current `DotProductAttention` when multiple such modules
 are concatenated, for instance in consecutive transformer blocks.
 apply_query_key_layer_scaling: bool, default = `False`
 apply query-key layer scaling during BMM1
 by a factor of `layer_number`
 attention_softmax_in_fp32: bool, default = `True`
 if set to `False`, softmax is executed in
 the dtype of activation tensors.
 attn_mask_type: {'causal', 'padding'}, default = `causal`
 type of attention mask passed into softmax operation.

Expand DownExpand Up@@ -371,9 +366,8 @@ def __init__(

 self.use_flash_attention = (
 int(os.getenv("NVTE_FLASH_ATTN", "1"))
 and attention_softmax_in_fp32
 and attn_mask_type == "causal"
 and not apply_query_key_layer_scaling
 and get_device_compute_capability() >= 8.0
 )

 attn_kwargs = {
Expand DownExpand Up@@ -422,6 +416,11 @@ def forward(
 """
 Dot Product Attention Layer.

 .. note::

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

 .. note::

 Input tensors :attr:`query_layer`, :attr:`key_layer`, and :attr:`value_layer`
Expand All@@ -448,8 +447,7 @@ def forward(
 """

 use_flash_attention = self.use_flash_attention
 if (attention_mask is not None
 or query_layer.dtype not in [torch.bfloat16, torch.float16]
 if (query_layer.dtype not in [torch.bfloat16, torch.float16]
 or key_layer.dtype not in [torch.bfloat16, torch.float16]
 or value_layer.dtype not in [torch.bfloat16, torch.float16]
 ):
Expand DownExpand Up@@ -515,6 +513,7 @@ def __init__(
 self.return_layernorm_output = return_layernorm_output
 self.params_dtype = params_dtype
 self.init_method = init_method
 self.attn_mask_type = attn_mask_type

 if not fuse_qkv_params:
 qkv_weight_interleaved = False
Expand DownExpand Up@@ -658,7 +657,7 @@ def forward(
 """MultiHeadAttention FWD"""
 # hidden_states: [sq, b, h]

 if attention_mask is not None:
 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"
Expand DownExpand Up@@ -836,6 +835,11 @@ class TransformerLayer(torch.nn.Module):
 TransformerLayer is made up of an attention block and a feedforward network (MLP).
 This standard layer is based on the paper "Attention Is All You Need".

 .. note::

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

 Parameters
 ----------
 hidden_size : int
Expand DownExpand Up@@ -983,6 +987,7 @@ def __init__(
 self.apply_residual_connection_post_layernorm = (
 apply_residual_connection_post_layernorm
 )
 self.self_attn_mask_type = self_attn_mask_type
 assert (
 self_attn_mask_type in AttnMaskTypes
 ), f"self_attn_mask_type {self_attn_mask_type} not supported"
Expand DownExpand Up@@ -1129,6 +1134,11 @@ def forward(
 """
 Transformer Layer: attention block and a feedforward network (MLP)

 .. note::

 Argument :attr:`attention_mask` will be ignored when :attr:`self_attn_mask_type`
 is set to `"causal"`.

 Parameters
 ----------
 hidden_states : torch.Tensor
Expand DownExpand Up@@ -1163,7 +1173,7 @@ def forward(

 hidden_states = hidden_states.contiguous()

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

7 changes: 7 additions & 0 deletions
transformer_engine/pytorch/utils.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -8,6 +8,13 @@
import torch


def get_device_compute_capability() -> float:
 """Returns the cuda compute capability of current GPU"""
 major = torch.cuda.get_device_properties(torch.cuda.current_device()).major
 minor = torch.cuda.get_device_properties(torch.cuda.current_device()).minor
 return major + minor / 10


def attention_mask_func(
 attention_scores: torch.Tensor, attention_mask: torch.Tensor
) -> torch.Tensor:
Expand Down

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

Skip to content
Fix flash attention#84

New issue

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

 Sign up for GitHub

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

Already on GitHub?
Sign in
to your account

 Merged

ksivaman
merged 5 commits into

from

Mar 7, 2023
 Merged


Fix flash attention
#84
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


5 commits


Select commit
Hold shift + click to select a range
f3184f3

ignore self attention mask for causal type


ksivaman Mar 2, 2023
20317c2

further relax checks to run FA, update docs


ksivaman Mar 2, 2023
b0ef5ab

fix pytorch softmax path


ksivaman Mar 2, 2023
23c5616

fixes


ksivaman Mar 6, 2023
f696393

minimum ampere requirement for fa


ksivaman Mar 7, 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


1 change: 1 addition & 0 deletions
tests/pytorch/test_onnx_export.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -792,6 +792,7 @@ def test_export_core_attention(

 if attn_mask_type is None:
 attn_mask_type = 'causal'
 inp = (query_layer, key_layer, value_layer)
 model = te.transformer.DotProductAttention(
 num_attention_heads=num_attention_heads,
 kv_channels=kv_channels,
Expand Down

13 changes: 13 additions & 0 deletions
transformer_engine/pytorch/softmax.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,6 +16,15 @@
THREADS_PER_BLOCK = 128


_default_causal_mask = {}

def _get_default_causal_mask(sq: int) -> torch.Tensor:
 """Return the causal upper triangular mask for softmax input"""
 if sq not in _default_causal_mask:
 _default_causal_mask[sq] = torch.triu(torch.ones(sq, sq, device="cuda"), diagonal=1).bool()
 return _default_causal_mask[sq]


class ScaledUpperTriangMaskedSoftmax(torch.autograd.Function):
 """
 Fused operation which performs following three operations in sequence
Expand DownExpand Up@@ -274,6 +283,10 @@ def forward_torch_softmax(

 if self.scale is not None:
 inp = inp * self.scale

 if self.attn_mask_type == "causal":
 mask = _get_default_causal_mask(inp.size()[2])

 mask_output = self.mask_func(inp, mask) if mask is not None else inp
 probs = torch.nn.Softmax(dim=-1)(mask_output)

Expand Down

46 changes: 28 additions & 18 deletions
transformer_engine/pytorch/transformer.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -27,6 +27,7 @@
 split_tensor_along_dim,
 cast_if_needed,
 get_default_init_method,
 get_device_compute_capability,
)
from transformer_engine.pytorch.constants import (
 AttnMaskTypes,
Expand DownExpand Up@@ -220,16 +221,14 @@ def __init__(
 assert (
 attn_mask_type == "causal"
 ), 'FlashAttention currently only supports causal attention mask.'
 assert (
 attention_softmax_in_fp32
 ), 'FlashAttention currently only supports softmax compute in fp32.'

 self.attn_causal_mask = attn_mask_type == "causal"
 self.norm_factor = norm_factor
 self.attention_dropout_ctx = attention_dropout_ctx
 self.attention_dropout = attention_dropout
 self.layer_number = layer_number
 self.apply_query_key_layer_scaling = apply_query_key_layer_scaling
 self.attention_softmax_in_fp32 = attention_softmax_in_fp32

 def forward(
 self,
Expand DownExpand Up@@ -287,6 +286,11 @@ class DotProductAttention(torch.nn.Module):
 representation subspaces as described in the paper:
 `Attention Is All You Need <https://arxiv.org/abs/1706.03762>`_.

 .. note::

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

 .. warning::

 For the default attention mechanism, this module executes a non-deterministic version of
Expand All@@ -303,15 +307,6 @@ 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.
 layer_number: int, default = `None`
 layer number of the current `DotProductAttention` when multiple such modules
 are concatenated, for instance in consecutive transformer blocks.
 apply_query_key_layer_scaling: bool, default = `False`
 apply query-key layer scaling during BMM1
 by a factor of `layer_number`
 attention_softmax_in_fp32: bool, default = `True`
 if set to `False`, softmax is executed in
 the dtype of activation tensors.
 attn_mask_type: {'causal', 'padding'}, default = `causal`
 type of attention mask passed into softmax operation.

Expand DownExpand Up@@ -371,9 +366,8 @@ def __init__(

 self.use_flash_attention = (
 int(os.getenv("NVTE_FLASH_ATTN", "1"))
 and attention_softmax_in_fp32
 and attn_mask_type == "causal"
 and not apply_query_key_layer_scaling
 and get_device_compute_capability() >= 8.0
 )

 attn_kwargs = {
Expand DownExpand Up@@ -422,6 +416,11 @@ def forward(
 """
 Dot Product Attention Layer.

 .. note::

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

 .. note::

 Input tensors :attr:`query_layer`, :attr:`key_layer`, and :attr:`value_layer`
Expand All@@ -448,8 +447,7 @@ def forward(
 """

 use_flash_attention = self.use_flash_attention
 if (attention_mask is not None
 or query_layer.dtype not in [torch.bfloat16, torch.float16]
 if (query_layer.dtype not in [torch.bfloat16, torch.float16]
 or key_layer.dtype not in [torch.bfloat16, torch.float16]
 or value_layer.dtype not in [torch.bfloat16, torch.float16]
 ):
Expand DownExpand Up@@ -515,6 +513,7 @@ def __init__(
 self.return_layernorm_output = return_layernorm_output
 self.params_dtype = params_dtype
 self.init_method = init_method
 self.attn_mask_type = attn_mask_type

 if not fuse_qkv_params:
 qkv_weight_interleaved = False
Expand DownExpand Up@@ -658,7 +657,7 @@ def forward(
 """MultiHeadAttention FWD"""
 # hidden_states: [sq, b, h]

 if attention_mask is not None:
 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"
Expand DownExpand Up@@ -836,6 +835,11 @@ class TransformerLayer(torch.nn.Module):
 TransformerLayer is made up of an attention block and a feedforward network (MLP).
 This standard layer is based on the paper "Attention Is All You Need".

 .. note::

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

 Parameters
 ----------
 hidden_size : int
Expand DownExpand Up@@ -983,6 +987,7 @@ def __init__(
 self.apply_residual_connection_post_layernorm = (
 apply_residual_connection_post_layernorm
 )
 self.self_attn_mask_type = self_attn_mask_type
 assert (
 self_attn_mask_type in AttnMaskTypes
 ), f"self_attn_mask_type {self_attn_mask_type} not supported"
Expand DownExpand Up@@ -1129,6 +1134,11 @@ def forward(
 """
 Transformer Layer: attention block and a feedforward network (MLP)

 .. note::

 Argument :attr:`attention_mask` will be ignored when :attr:`self_attn_mask_type`
 is set to `"causal"`.

 Parameters
 ----------
 hidden_states : torch.Tensor
Expand DownExpand Up@@ -1163,7 +1173,7 @@ def forward(

 hidden_states = hidden_states.contiguous()

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

7 changes: 7 additions & 0 deletions
transformer_engine/pytorch/utils.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -8,6 +8,13 @@
import torch


def get_device_compute_capability() -> float:
 """Returns the cuda compute capability of current GPU"""
 major = torch.cuda.get_device_properties(torch.cuda.current_device()).major
 minor = torch.cuda.get_device_properties(torch.cuda.current_device()).minor
 return major + minor / 10


def attention_mask_func(
 attention_scores: torch.Tensor, attention_mask: torch.Tensor
) -> torch.Tensor:
Expand Down