Skip to content

flash-attn integration - #62

Merged
ksivaman merged 39 commits into
NVIDIA:mainfrom
cyanguwa:flash_attn
Feb 22, 2023
Merged

flash-attn integration#62
ksivaman merged 39 commits into
NVIDIA:mainfrom
cyanguwa:flash_attn

Conversation

@cyanguwa

@cyanguwacyanguwa commented Jan 31, 2023

Copy link
Copy Markdown
Collaborator

This PR is to provide support for flash attention in TE for self-attention calculations.

@ksivaman
ksivaman marked this pull request as draft January 31, 2023 00:53
@ksivaman

ksivaman commented Feb 4, 2023

Copy link
Copy Markdown
Member

@cyanguwa Could you sign off the commits? This guide explains how to do it.

@ksivaman
ksivaman marked this pull request as ready for review February 13, 2023 19:16
@ksivaman
ksivaman self-requested a review February 13, 2023 19:17
cyanguwaand others added 13 commits February 13, 2023 14:31
@cyanguwa
add flash attention to TransformerLayer
75ecfff
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Add docs for FP8 calibration (NVIDIA#61)
f428fe4
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ptrendx@cyanguwa
Fix the integer overflow in fused softmax (NVIDIA#60)
9c67222
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
prefix flash attn env var with NVTE_
faa7f5b
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Address steady memory increase and bloated checkpoints (NVIDIA#63)
ecef587
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
fix env var logic
1dccf39
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
fix flash attn env var logic again
38168dc
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
remove d2d copies (NVIDIA#64)
6ad62d3
* remove d2d copies
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* cleanup
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@vasunvidia@cyanguwa
Increase number of FP8 tensors per GEMM (NVIDIA#22)
f071530
* Increase number of FP8 tensors per GEMM
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Enable FP8 output tensor for fp8_gemm
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* [BERT FP8] Initial TE review comments
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Temporary fix for cuda graph non convergence
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Address review comments-2
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Review comments-3
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Cleanup
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Change for New API
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Remove unnecessary clone for D_scale, D_amax
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Avoid Roll for AMAX history size = 1
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Update onnx_te_gemm API
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Fix Lint errors
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
---------
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Bug fixes from PR 22 (NVIDIA#65)
d79c1b4
* Bug fixes from PR 22
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Add FP8 tests to ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* bundle unittests for ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
replace rearrange with transpose
be59386
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
QKV parameters unfused path fixes and optimization (NVIDIA#66)
faf9146
* Bug fixes from PR 22
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Add FP8 tests to ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Better QKV parameter fusion
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* small fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* keep original param for unfused case to retain externally set attrs
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* lint fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Fix ONNX exports
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* improve arg naming
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* No need to set data pointers
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* lint
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Assert memory loc in NoopCat
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Handle case of different memory in param and buffer
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* fix assert always true
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Reassign params memory to avoid more concats
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Fix gradients when using AMP (NVIDIA#70)
2416cc7
retain grad related attrs while casting
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
cyanguwaand others added 8 commits February 13, 2023 14:42
@cyanguwa
Merge branch 'main' into flash_attn
31cd706
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
fix pylint violations
d0f2eaf
fixed pyline violations such as trailing white spaces and too long lines Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
fix pylint violation on line 264 with R1719
304ef79
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
fix two more pylint violations
b3440c2
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
Merge branch 'NVIDIA:main' into flash_attn
cf28311
@ksivaman
DotProductAttention API
0b34943
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
Add docs for attention
39d4d71
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
fix assert always true
bd5160c
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman

ksivaman commented Feb 16, 2023

Copy link
Copy Markdown
Member

/te-ci

@ksivaman
check for correct flash-attn version
4d8a6cb
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman

ksivaman commented Feb 16, 2023

Copy link
Copy Markdown
Member

/te-ci

@ksivaman
review comments and fixes
28b571c
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
fix onnx and disable flash-attn export test
9b6b6c6
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
remove einops dependency
0551b43
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
cleanup internal API; rm duplication
cd6c90c
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
fix
9f82f42
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
only install TE wheel (exclude flash-attn to rm conflicts)
c9783b7
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman

ksivaman commented Feb 17, 2023

Copy link
Copy Markdown
Member

/te-ci

@ksivaman
forgot to change install wheel path
ed64331
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
ptrendx
ptrendx reviewed Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py Outdated
softmax_scale=1.0/self.norm_factor, causal=self.attn_causal_mask
)

# [b, sq, np, hn]

@ptrendxptrendxFeb 17, 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.

Wrong comment or wrong code?

@ksivamanksivamanFeb 17, 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.

The dimension comments are to indicate shape before not after

@ptrendxptrendxFeb 17, 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.

Well, then the code is wrong, since to get that you would need to do transpose, not view, right?

@ksivamanksivamanFeb 18, 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.

You are actually right, everything is wrong here; comment (input shape) and the view. I couldn't catch in my numerical tests since I'm running with mbs=1, sigh

ptrendx
ptrendx reviewed Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py Outdated
ptrendx
ptrendx reviewed Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
ptrendx
ptrendx reviewed Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py Outdated
@ksivaman
next round review comments
09687cc
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman

ksivaman commented Feb 17, 2023

Copy link
Copy Markdown
Member

/te-ci

@ksivaman
fix flash_attn output
c39d043
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
ptrendx
ptrendx approved these changes Feb 19, 2023
View reviewed changes

@ptrendxptrendx left a comment

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.

LGTM, thanks!

@ksivaman
fix QK layer scaling
1da83df
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
update docs
dc51968
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman

ksivaman commented Feb 22, 2023

Copy link
Copy Markdown
Member

/te-ci

@ksivaman

ksivaman commented Feb 22, 2023

Copy link
Copy Markdown
Member

ci error is unrelated ONNX failure due to tolerance

cyanguwa
cyanguwa commented Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py Outdated
), 'FlashAttention currently only supports CUDA tensors.'
assert (
attention_mask is None
), 'FlashAttention currently does not support attention mask.'

@cyanguwacyanguwaFeb 22, 2023

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

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

FA calculates the causal attention mask inside itself so maybe here we meant 'FlashAttention doesn't support external attention mask'?

cyanguwa
cyanguwa commented Feb 22, 2023
View reviewed changes
Comment threadqa/L0_unittest/test.sh
pip install pytest==6.2.5 onnxruntime==1.13.1
pytest -v -s $TE_PATH/tests/*.py
pytest -v -s $TE_PATH/tests/test_transformerengine.py $TE_PATH/tests/test_fp8.py
NVTE_FLASH_ATTN=0 pytest -v -s $TE_PATH/tests/test_onnx_export.py

@cyanguwacyanguwaFeb 22, 2023

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

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

We don't test NVTE_FLASH_ATTN=1 I guess for ONNX?

cyanguwa
cyanguwa commented Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
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

@cyanguwacyanguwaFeb 22, 2023

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

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

Do we need L231-232 since we don't use layer_number and apply_query_key_layer_scaling in FA? I guess you're trying to keep the interface consistent between Unfused and Flash?

@ksivamanksivamanFeb 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.

Yes the second point is correct, just for consistency

cyanguwa
cyanguwa commented Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py Outdated

if use_flash_attention:
if checkpoint_core_attention:
return self._checkpointed_attention_forward(query_layer, key_layer, value_layer)

@cyanguwacyanguwaFeb 22, 2023

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

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

Did we forget to pass in the attention_func here? Same for L457.

@ksivamanksivamanFeb 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.

Yes there are some bugs in the selective activation checkpointing due to the MLM integration that doesn't actually run this path when I test for it, thus I wrongly assumed it works, fixing it now. Also will add standalone tests for this separately. Good catch :)

cyanguwa
cyanguwa commented Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
custom_forward,
False,
self.get_rng_state_tracker,
self.tp_group,

@cyanguwacyanguwaFeb 22, 2023

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

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

Does self have a tp_group member?

@ksivamanksivamanFeb 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.

Same as above

@ksivaman
review comments and fixes to selective checkpointing
14a2beb
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman

ksivaman commented Feb 22, 2023

Copy link
Copy Markdown
Member

/te-ci

@ksivaman
ksivaman merged commit 79a9fe2 into NVIDIA:mainFeb 22, 2023
@jeng1220jeng1220 mentioned this pull request Feb 22, 2023
Merged
@ksivamanksivaman mentioned this pull request Feb 24, 2023
Closed
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ksivamanksivamanksivaman left review comments
@ptrendxptrendxptrendx approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@cyanguwa@ksivaman@ptrendx@vasunvidia
, '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" + '
flash-attn integration by cyanguwa · Pull Request #62 · NVIDIA/TransformerEngine · GitHub
Skip to content
flash-attn integration - #62
Merged
ksivaman merged 39 commits into
NVIDIA:mainfrom 
cyanguwa:flash_attn
Feb 22, 2023
Merged
flash-attn integration#62
ksivaman merged 39 commits into
NVIDIA:mainfrom 
cyanguwa:flash_attn
Conversation
@cyanguwa
@cyanguwacyanguwa
commented
Jan 31, 2023

Copy link


Copy Markdown

Collaborator
This PR is to provide support for flash attention in TE for self-attention calculations.
@ksivaman
ksivaman
marked this pull request as draft
January 31, 2023 00:53
@ksivaman
ksivaman
commented
Feb 4, 2023

Copy link


Copy Markdown

Member
@cyanguwa Could you sign off the commits? This guide explains how to do it.
@ksivaman
ksivaman
marked this pull request as ready for review
February 13, 2023 19:16
@ksivaman
ksivaman
self-requested a review
February 13, 2023 19:17
cyanguwaand others
added 13 commits
February 13, 2023 14:31
@cyanguwa
add flash attention to TransformerLayer
75ecfff
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Add docs for FP8 calibration (NVIDIA#61)
f428fe4
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ptrendx@cyanguwa
Fix the integer overflow in fused softmax (NVIDIA#60)
9c67222
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
prefix flash attn env var with NVTE_
faa7f5b
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Address steady memory increase and bloated checkpoints (NVIDIA#63)
ecef587
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
fix env var logic
1dccf39
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
fix flash attn env var logic again
38168dc
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
remove d2d copies (NVIDIA#64)
6ad62d3
* remove d2d copies
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* cleanup
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@vasunvidia@cyanguwa
Increase number of FP8 tensors per GEMM (NVIDIA#22)
f071530
* Increase number of FP8 tensors per GEMM
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Enable FP8 output tensor for fp8_gemm
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* [BERT FP8] Initial TE review comments
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Temporary fix for cuda graph non convergence
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Address review comments-2
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Review comments-3
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Cleanup
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Change for New API
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Remove unnecessary clone for D_scale, D_amax
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Avoid Roll for AMAX history size = 1
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Update onnx_te_gemm API
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Fix Lint errors
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
---------
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Bug fixes from PR 22 (NVIDIA#65)
d79c1b4
* Bug fixes from PR 22
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Add FP8 tests to ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* bundle unittests for ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
replace rearrange with transpose
be59386
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
QKV parameters unfused path fixes and optimization (NVIDIA#66)
faf9146
* Bug fixes from PR 22
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Add FP8 tests to ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Better QKV parameter fusion
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* small fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* keep original param for unfused case to retain externally set attrs
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* lint fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Fix ONNX exports
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* improve arg naming
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* No need to set data pointers
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* lint
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Assert memory loc in NoopCat
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Handle case of different memory in param and buffer
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* fix assert always true
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Reassign params memory to avoid more concats
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Fix gradients when using AMP (NVIDIA#70)
2416cc7
retain grad related attrs while casting
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
cyanguwaand others
added 8 commits
February 13, 2023 14:42
@cyanguwa
Merge branch 'main' into flash_attn
31cd706
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
fix pylint violations
d0f2eaf
fixed pyline violations such as trailing white spaces and too long lines Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
fix pylint violation on line 264 with R1719
304ef79
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
fix two more pylint violations
b3440c2
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
Merge branch 'NVIDIA:main' into flash_attn
cf28311
@ksivaman
DotProductAttention API
0b34943
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
Add docs for attention
39d4d71
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
fix assert always true
bd5160c
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 16, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
check for correct flash-attn version
4d8a6cb
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 16, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
review comments and fixes
28b571c
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
fix onnx and disable flash-attn export test
9b6b6c6
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
remove einops dependency
0551b43
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
cleanup internal API; rm duplication
cd6c90c
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
fix
9f82f42
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
only install TE wheel (exclude flash-attn to rm conflicts)
c9783b7
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 17, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
forgot to change install wheel path
ed64331
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
ptrendx
ptrendx
reviewed
Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated

 softmax_scale=1.0/self.norm_factor, causal=self.attn_causal_mask
 )

 # [b, sq, np, hn]
@ptrendxptrendxFeb 17, 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.

Wrong comment or wrong code?
@ksivamanksivamanFeb 17, 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.

The dimension comments are to indicate shape before not after
@ptrendxptrendxFeb 17, 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.

Well, then the code is wrong, since to get that you would need to do transpose, not view, right?
@ksivamanksivamanFeb 18, 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.

You are actually right, everything is wrong here; comment (input shape) and the view. I couldn't catch in my numerical tests since I'm running with mbs=1, sigh
ptrendx
ptrendx
reviewed
Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated

ptrendx
ptrendx
reviewed
Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
ptrendx
ptrendx
reviewed
Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated

@ksivaman
next round review comments
09687cc
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 17, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
fix flash_attn output
c39d043
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
ptrendx
ptrendx
approved these changes
Feb 19, 2023
View reviewed changes
@ptrendxptrendx
left a comment


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.

LGTM, thanks!
@ksivaman
fix QK layer scaling
1da83df
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
update docs
dc51968
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 22, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
ksivaman
commented
Feb 22, 2023

Copy link


Copy Markdown

Member
ci error is unrelated ONNX failure due to tolerance
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated

 ), 'FlashAttention currently only supports CUDA tensors.'
 assert (
 attention_mask is None
 ), 'FlashAttention currently does not support attention mask.'
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

FA calculates the causal attention mask inside itself so maybe here we meant 'FlashAttention doesn't support external attention mask'?
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadqa/L0_unittest/test.sh
pip install pytest==6.2.5 onnxruntime==1.13.1
pytest -v -s $TE_PATH/tests/*.py
pytest -v -s $TE_PATH/tests/test_transformerengine.py $TE_PATH/tests/test_fp8.py
NVTE_FLASH_ATTN=0 pytest -v -s $TE_PATH/tests/test_onnx_export.py
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

We don't test NVTE_FLASH_ATTN=1 I guess for ONNX?
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
 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
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Do we need L231-232 since we don't use layer_number and apply_query_key_layer_scaling in FA? I guess you're trying to keep the interface consistent between Unfused and Flash?
@ksivamanksivamanFeb 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.

Yes the second point is correct, just for consistency
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated


 if use_flash_attention:
 if checkpoint_core_attention:
 return self._checkpointed_attention_forward(query_layer, key_layer, value_layer)
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Did we forget to pass in the attention_func here? Same for L457.
@ksivamanksivamanFeb 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.

Yes there are some bugs in the selective activation checkpointing due to the MLM integration that doesn't actually run this path when I test for it, thus I wrongly assumed it works, fixing it now. Also will add standalone tests for this separately. Good catch :)
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
 custom_forward,
 False,
 self.get_rng_state_tracker,
 self.tp_group,
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Does self have a tp_group member?
@ksivamanksivamanFeb 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.

Same as above
@ksivaman
review comments and fixes to selective checkpointing
14a2beb
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 22, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
ksivaman
merged commit 79a9fe2
into
NVIDIA:mainFeb 22, 2023
@jeng1220jeng1220
mentioned this pull request
Feb 22, 2023
 Merged

@ksivamanksivaman
mentioned this pull request
Feb 24, 2023
 Closed

Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@ksivamanksivamanksivaman left review comments
@ptrendxptrendxptrendx approved these changes

Assignees


No one assigned


Labels


None yet


Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

4 participants

@cyanguwa@ksivaman@ptrendx@vasunvidia

, '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('^' + ".*" + '
flash-attn integration by cyanguwa · Pull Request #62 · NVIDIA/TransformerEngine · GitHub
Skip to content
flash-attn integration - #62
Merged
ksivaman merged 39 commits into
NVIDIA:mainfrom 
cyanguwa:flash_attn
Feb 22, 2023
Merged
flash-attn integration#62
ksivaman merged 39 commits into
NVIDIA:mainfrom 
cyanguwa:flash_attn
Conversation
@cyanguwa
@cyanguwacyanguwa
commented
Jan 31, 2023

Copy link


Copy Markdown

Collaborator
This PR is to provide support for flash attention in TE for self-attention calculations.
@ksivaman
ksivaman
marked this pull request as draft
January 31, 2023 00:53
@ksivaman
ksivaman
commented
Feb 4, 2023

Copy link


Copy Markdown

Member
@cyanguwa Could you sign off the commits? This guide explains how to do it.
@ksivaman
ksivaman
marked this pull request as ready for review
February 13, 2023 19:16
@ksivaman
ksivaman
self-requested a review
February 13, 2023 19:17
cyanguwaand others
added 13 commits
February 13, 2023 14:31
@cyanguwa
add flash attention to TransformerLayer
75ecfff
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Add docs for FP8 calibration (NVIDIA#61)
f428fe4
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ptrendx@cyanguwa
Fix the integer overflow in fused softmax (NVIDIA#60)
9c67222
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
prefix flash attn env var with NVTE_
faa7f5b
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Address steady memory increase and bloated checkpoints (NVIDIA#63)
ecef587
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
fix env var logic
1dccf39
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
fix flash attn env var logic again
38168dc
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
remove d2d copies (NVIDIA#64)
6ad62d3
* remove d2d copies
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* cleanup
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@vasunvidia@cyanguwa
Increase number of FP8 tensors per GEMM (NVIDIA#22)
f071530
* Increase number of FP8 tensors per GEMM
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Enable FP8 output tensor for fp8_gemm
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* [BERT FP8] Initial TE review comments
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Temporary fix for cuda graph non convergence
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Address review comments-2
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Review comments-3
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Cleanup
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Change for New API
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Remove unnecessary clone for D_scale, D_amax
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Avoid Roll for AMAX history size = 1
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Update onnx_te_gemm API
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Fix Lint errors
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
---------
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Bug fixes from PR 22 (NVIDIA#65)
d79c1b4
* Bug fixes from PR 22
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Add FP8 tests to ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* bundle unittests for ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
replace rearrange with transpose
be59386
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
QKV parameters unfused path fixes and optimization (NVIDIA#66)
faf9146
* Bug fixes from PR 22
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Add FP8 tests to ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Better QKV parameter fusion
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* small fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* keep original param for unfused case to retain externally set attrs
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* lint fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Fix ONNX exports
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* improve arg naming
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* No need to set data pointers
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* lint
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Assert memory loc in NoopCat
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Handle case of different memory in param and buffer
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* fix assert always true
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Reassign params memory to avoid more concats
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Fix gradients when using AMP (NVIDIA#70)
2416cc7
retain grad related attrs while casting
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
cyanguwaand others
added 8 commits
February 13, 2023 14:42
@cyanguwa
Merge branch 'main' into flash_attn
31cd706
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
fix pylint violations
d0f2eaf
fixed pyline violations such as trailing white spaces and too long lines Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
fix pylint violation on line 264 with R1719
304ef79
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
fix two more pylint violations
b3440c2
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
Merge branch 'NVIDIA:main' into flash_attn
cf28311
@ksivaman
DotProductAttention API
0b34943
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
Add docs for attention
39d4d71
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
fix assert always true
bd5160c
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 16, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
check for correct flash-attn version
4d8a6cb
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 16, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
review comments and fixes
28b571c
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
fix onnx and disable flash-attn export test
9b6b6c6
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
remove einops dependency
0551b43
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
cleanup internal API; rm duplication
cd6c90c
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
fix
9f82f42
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
only install TE wheel (exclude flash-attn to rm conflicts)
c9783b7
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 17, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
forgot to change install wheel path
ed64331
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
ptrendx
ptrendx
reviewed
Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated

 softmax_scale=1.0/self.norm_factor, causal=self.attn_causal_mask
 )

 # [b, sq, np, hn]
@ptrendxptrendxFeb 17, 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.

Wrong comment or wrong code?
@ksivamanksivamanFeb 17, 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.

The dimension comments are to indicate shape before not after
@ptrendxptrendxFeb 17, 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.

Well, then the code is wrong, since to get that you would need to do transpose, not view, right?
@ksivamanksivamanFeb 18, 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.

You are actually right, everything is wrong here; comment (input shape) and the view. I couldn't catch in my numerical tests since I'm running with mbs=1, sigh
ptrendx
ptrendx
reviewed
Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated

ptrendx
ptrendx
reviewed
Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
ptrendx
ptrendx
reviewed
Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated

@ksivaman
next round review comments
09687cc
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 17, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
fix flash_attn output
c39d043
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
ptrendx
ptrendx
approved these changes
Feb 19, 2023
View reviewed changes
@ptrendxptrendx
left a comment


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.

LGTM, thanks!
@ksivaman
fix QK layer scaling
1da83df
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
update docs
dc51968
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 22, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
ksivaman
commented
Feb 22, 2023

Copy link


Copy Markdown

Member
ci error is unrelated ONNX failure due to tolerance
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated

 ), 'FlashAttention currently only supports CUDA tensors.'
 assert (
 attention_mask is None
 ), 'FlashAttention currently does not support attention mask.'
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

FA calculates the causal attention mask inside itself so maybe here we meant 'FlashAttention doesn't support external attention mask'?
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadqa/L0_unittest/test.sh
pip install pytest==6.2.5 onnxruntime==1.13.1
pytest -v -s $TE_PATH/tests/*.py
pytest -v -s $TE_PATH/tests/test_transformerengine.py $TE_PATH/tests/test_fp8.py
NVTE_FLASH_ATTN=0 pytest -v -s $TE_PATH/tests/test_onnx_export.py
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

We don't test NVTE_FLASH_ATTN=1 I guess for ONNX?
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
 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
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Do we need L231-232 since we don't use layer_number and apply_query_key_layer_scaling in FA? I guess you're trying to keep the interface consistent between Unfused and Flash?
@ksivamanksivamanFeb 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.

Yes the second point is correct, just for consistency
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated


 if use_flash_attention:
 if checkpoint_core_attention:
 return self._checkpointed_attention_forward(query_layer, key_layer, value_layer)
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Did we forget to pass in the attention_func here? Same for L457.
@ksivamanksivamanFeb 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.

Yes there are some bugs in the selective activation checkpointing due to the MLM integration that doesn't actually run this path when I test for it, thus I wrongly assumed it works, fixing it now. Also will add standalone tests for this separately. Good catch :)
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
 custom_forward,
 False,
 self.get_rng_state_tracker,
 self.tp_group,
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Does self have a tp_group member?
@ksivamanksivamanFeb 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.

Same as above
@ksivaman
review comments and fixes to selective checkpointing
14a2beb
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 22, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
ksivaman
merged commit 79a9fe2
into
NVIDIA:mainFeb 22, 2023
@jeng1220jeng1220
mentioned this pull request
Feb 22, 2023
 Merged

@ksivamanksivaman
mentioned this pull request
Feb 24, 2023
 Closed

Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@ksivamanksivamanksivaman left review comments
@ptrendxptrendxptrendx approved these changes

Assignees


No one assigned


Labels


None yet


Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

4 participants

@cyanguwa@ksivaman@ptrendx@vasunvidia

, '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('^' + ".*" + '
flash-attn integration by cyanguwa · Pull Request #62 · NVIDIA/TransformerEngine · GitHub
Skip to content
flash-attn integration - #62
Merged
ksivaman merged 39 commits into
NVIDIA:mainfrom 
cyanguwa:flash_attn
Feb 22, 2023
Merged
flash-attn integration#62
ksivaman merged 39 commits into
NVIDIA:mainfrom 
cyanguwa:flash_attn
Conversation
@cyanguwa
@cyanguwacyanguwa
commented
Jan 31, 2023

Copy link


Copy Markdown

Collaborator
This PR is to provide support for flash attention in TE for self-attention calculations.
@ksivaman
ksivaman
marked this pull request as draft
January 31, 2023 00:53
@ksivaman
ksivaman
commented
Feb 4, 2023

Copy link


Copy Markdown

Member
@cyanguwa Could you sign off the commits? This guide explains how to do it.
@ksivaman
ksivaman
marked this pull request as ready for review
February 13, 2023 19:16
@ksivaman
ksivaman
self-requested a review
February 13, 2023 19:17
cyanguwaand others
added 13 commits
February 13, 2023 14:31
@cyanguwa
add flash attention to TransformerLayer
75ecfff
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Add docs for FP8 calibration (NVIDIA#61)
f428fe4
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ptrendx@cyanguwa
Fix the integer overflow in fused softmax (NVIDIA#60)
9c67222
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
prefix flash attn env var with NVTE_
faa7f5b
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Address steady memory increase and bloated checkpoints (NVIDIA#63)
ecef587
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
fix env var logic
1dccf39
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
fix flash attn env var logic again
38168dc
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
remove d2d copies (NVIDIA#64)
6ad62d3
* remove d2d copies
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* cleanup
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@vasunvidia@cyanguwa
Increase number of FP8 tensors per GEMM (NVIDIA#22)
f071530
* Increase number of FP8 tensors per GEMM
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Enable FP8 output tensor for fp8_gemm
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* [BERT FP8] Initial TE review comments
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Temporary fix for cuda graph non convergence
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Address review comments-2
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Review comments-3
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Cleanup
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Change for New API
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Remove unnecessary clone for D_scale, D_amax
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Avoid Roll for AMAX history size = 1
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Update onnx_te_gemm API
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Fix Lint errors
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
---------
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Bug fixes from PR 22 (NVIDIA#65)
d79c1b4
* Bug fixes from PR 22
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Add FP8 tests to ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* bundle unittests for ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
replace rearrange with transpose
be59386
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
QKV parameters unfused path fixes and optimization (NVIDIA#66)
faf9146
* Bug fixes from PR 22
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Add FP8 tests to ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Better QKV parameter fusion
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* small fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* keep original param for unfused case to retain externally set attrs
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* lint fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Fix ONNX exports
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* improve arg naming
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* No need to set data pointers
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* lint
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Assert memory loc in NoopCat
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Handle case of different memory in param and buffer
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* fix assert always true
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Reassign params memory to avoid more concats
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Fix gradients when using AMP (NVIDIA#70)
2416cc7
retain grad related attrs while casting
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
cyanguwaand others
added 8 commits
February 13, 2023 14:42
@cyanguwa
Merge branch 'main' into flash_attn
31cd706
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
fix pylint violations
d0f2eaf
fixed pyline violations such as trailing white spaces and too long lines Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
fix pylint violation on line 264 with R1719
304ef79
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
fix two more pylint violations
b3440c2
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
Merge branch 'NVIDIA:main' into flash_attn
cf28311
@ksivaman
DotProductAttention API
0b34943
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
Add docs for attention
39d4d71
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
fix assert always true
bd5160c
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 16, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
check for correct flash-attn version
4d8a6cb
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 16, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
review comments and fixes
28b571c
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
fix onnx and disable flash-attn export test
9b6b6c6
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
remove einops dependency
0551b43
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
cleanup internal API; rm duplication
cd6c90c
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
fix
9f82f42
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
only install TE wheel (exclude flash-attn to rm conflicts)
c9783b7
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 17, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
forgot to change install wheel path
ed64331
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
ptrendx
ptrendx
reviewed
Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated

 softmax_scale=1.0/self.norm_factor, causal=self.attn_causal_mask
 )

 # [b, sq, np, hn]
@ptrendxptrendxFeb 17, 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.

Wrong comment or wrong code?
@ksivamanksivamanFeb 17, 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.

The dimension comments are to indicate shape before not after
@ptrendxptrendxFeb 17, 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.

Well, then the code is wrong, since to get that you would need to do transpose, not view, right?
@ksivamanksivamanFeb 18, 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.

You are actually right, everything is wrong here; comment (input shape) and the view. I couldn't catch in my numerical tests since I'm running with mbs=1, sigh
ptrendx
ptrendx
reviewed
Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated

ptrendx
ptrendx
reviewed
Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
ptrendx
ptrendx
reviewed
Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated

@ksivaman
next round review comments
09687cc
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 17, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
fix flash_attn output
c39d043
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
ptrendx
ptrendx
approved these changes
Feb 19, 2023
View reviewed changes
@ptrendxptrendx
left a comment


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.

LGTM, thanks!
@ksivaman
fix QK layer scaling
1da83df
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
update docs
dc51968
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 22, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
ksivaman
commented
Feb 22, 2023

Copy link


Copy Markdown

Member
ci error is unrelated ONNX failure due to tolerance
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated

 ), 'FlashAttention currently only supports CUDA tensors.'
 assert (
 attention_mask is None
 ), 'FlashAttention currently does not support attention mask.'
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

FA calculates the causal attention mask inside itself so maybe here we meant 'FlashAttention doesn't support external attention mask'?
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadqa/L0_unittest/test.sh
pip install pytest==6.2.5 onnxruntime==1.13.1
pytest -v -s $TE_PATH/tests/*.py
pytest -v -s $TE_PATH/tests/test_transformerengine.py $TE_PATH/tests/test_fp8.py
NVTE_FLASH_ATTN=0 pytest -v -s $TE_PATH/tests/test_onnx_export.py
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

We don't test NVTE_FLASH_ATTN=1 I guess for ONNX?
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
 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
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Do we need L231-232 since we don't use layer_number and apply_query_key_layer_scaling in FA? I guess you're trying to keep the interface consistent between Unfused and Flash?
@ksivamanksivamanFeb 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.

Yes the second point is correct, just for consistency
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated


 if use_flash_attention:
 if checkpoint_core_attention:
 return self._checkpointed_attention_forward(query_layer, key_layer, value_layer)
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Did we forget to pass in the attention_func here? Same for L457.
@ksivamanksivamanFeb 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.

Yes there are some bugs in the selective activation checkpointing due to the MLM integration that doesn't actually run this path when I test for it, thus I wrongly assumed it works, fixing it now. Also will add standalone tests for this separately. Good catch :)
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
 custom_forward,
 False,
 self.get_rng_state_tracker,
 self.tp_group,
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Does self have a tp_group member?
@ksivamanksivamanFeb 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.

Same as above
@ksivaman
review comments and fixes to selective checkpointing
14a2beb
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 22, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
ksivaman
merged commit 79a9fe2
into
NVIDIA:mainFeb 22, 2023
@jeng1220jeng1220
mentioned this pull request
Feb 22, 2023
 Merged

@ksivamanksivaman
mentioned this pull request
Feb 24, 2023
 Closed

Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@ksivamanksivamanksivaman left review comments
@ptrendxptrendxptrendx approved these changes

Assignees


No one assigned


Labels


None yet


Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

4 participants

@cyanguwa@ksivaman@ptrendx@vasunvidia

, '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" + '
flash-attn integration by cyanguwa · Pull Request #62 · NVIDIA/TransformerEngine · GitHub
Skip to content
flash-attn integration - #62
Merged
ksivaman merged 39 commits into
NVIDIA:mainfrom 
cyanguwa:flash_attn
Feb 22, 2023
Merged
flash-attn integration#62
ksivaman merged 39 commits into
NVIDIA:mainfrom 
cyanguwa:flash_attn
Conversation
@cyanguwa
@cyanguwacyanguwa
commented
Jan 31, 2023

Copy link


Copy Markdown

Collaborator
This PR is to provide support for flash attention in TE for self-attention calculations.
@ksivaman
ksivaman
marked this pull request as draft
January 31, 2023 00:53
@ksivaman
ksivaman
commented
Feb 4, 2023

Copy link


Copy Markdown

Member
@cyanguwa Could you sign off the commits? This guide explains how to do it.
@ksivaman
ksivaman
marked this pull request as ready for review
February 13, 2023 19:16
@ksivaman
ksivaman
self-requested a review
February 13, 2023 19:17
cyanguwaand others
added 13 commits
February 13, 2023 14:31
@cyanguwa
add flash attention to TransformerLayer
75ecfff
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Add docs for FP8 calibration (NVIDIA#61)
f428fe4
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ptrendx@cyanguwa
Fix the integer overflow in fused softmax (NVIDIA#60)
9c67222
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
prefix flash attn env var with NVTE_
faa7f5b
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Address steady memory increase and bloated checkpoints (NVIDIA#63)
ecef587
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
fix env var logic
1dccf39
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
fix flash attn env var logic again
38168dc
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
remove d2d copies (NVIDIA#64)
6ad62d3
* remove d2d copies
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* cleanup
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@vasunvidia@cyanguwa
Increase number of FP8 tensors per GEMM (NVIDIA#22)
f071530
* Increase number of FP8 tensors per GEMM
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Enable FP8 output tensor for fp8_gemm
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* [BERT FP8] Initial TE review comments
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Temporary fix for cuda graph non convergence
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Address review comments-2
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Review comments-3
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Cleanup
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Change for New API
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Remove unnecessary clone for D_scale, D_amax
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Avoid Roll for AMAX history size = 1
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Update onnx_te_gemm API
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Fix Lint errors
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
---------
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Bug fixes from PR 22 (NVIDIA#65)
d79c1b4
* Bug fixes from PR 22
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Add FP8 tests to ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* bundle unittests for ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
replace rearrange with transpose
be59386
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
QKV parameters unfused path fixes and optimization (NVIDIA#66)
faf9146
* Bug fixes from PR 22
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Add FP8 tests to ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Better QKV parameter fusion
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* small fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* keep original param for unfused case to retain externally set attrs
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* lint fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Fix ONNX exports
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* improve arg naming
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* No need to set data pointers
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* lint
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Assert memory loc in NoopCat
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Handle case of different memory in param and buffer
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* fix assert always true
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Reassign params memory to avoid more concats
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Fix gradients when using AMP (NVIDIA#70)
2416cc7
retain grad related attrs while casting
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
cyanguwaand others
added 8 commits
February 13, 2023 14:42
@cyanguwa
Merge branch 'main' into flash_attn
31cd706
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
fix pylint violations
d0f2eaf
fixed pyline violations such as trailing white spaces and too long lines Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
fix pylint violation on line 264 with R1719
304ef79
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
fix two more pylint violations
b3440c2
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
Merge branch 'NVIDIA:main' into flash_attn
cf28311
@ksivaman
DotProductAttention API
0b34943
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
Add docs for attention
39d4d71
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
fix assert always true
bd5160c
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 16, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
check for correct flash-attn version
4d8a6cb
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 16, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
review comments and fixes
28b571c
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
fix onnx and disable flash-attn export test
9b6b6c6
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
remove einops dependency
0551b43
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
cleanup internal API; rm duplication
cd6c90c
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
fix
9f82f42
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
only install TE wheel (exclude flash-attn to rm conflicts)
c9783b7
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 17, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
forgot to change install wheel path
ed64331
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
ptrendx
ptrendx
reviewed
Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated

 softmax_scale=1.0/self.norm_factor, causal=self.attn_causal_mask
 )

 # [b, sq, np, hn]
@ptrendxptrendxFeb 17, 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.

Wrong comment or wrong code?
@ksivamanksivamanFeb 17, 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.

The dimension comments are to indicate shape before not after
@ptrendxptrendxFeb 17, 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.

Well, then the code is wrong, since to get that you would need to do transpose, not view, right?
@ksivamanksivamanFeb 18, 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.

You are actually right, everything is wrong here; comment (input shape) and the view. I couldn't catch in my numerical tests since I'm running with mbs=1, sigh
ptrendx
ptrendx
reviewed
Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated

ptrendx
ptrendx
reviewed
Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
ptrendx
ptrendx
reviewed
Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated

@ksivaman
next round review comments
09687cc
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 17, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
fix flash_attn output
c39d043
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
ptrendx
ptrendx
approved these changes
Feb 19, 2023
View reviewed changes
@ptrendxptrendx
left a comment


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.

LGTM, thanks!
@ksivaman
fix QK layer scaling
1da83df
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
update docs
dc51968
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 22, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
ksivaman
commented
Feb 22, 2023

Copy link


Copy Markdown

Member
ci error is unrelated ONNX failure due to tolerance
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated

 ), 'FlashAttention currently only supports CUDA tensors.'
 assert (
 attention_mask is None
 ), 'FlashAttention currently does not support attention mask.'
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

FA calculates the causal attention mask inside itself so maybe here we meant 'FlashAttention doesn't support external attention mask'?
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadqa/L0_unittest/test.sh
pip install pytest==6.2.5 onnxruntime==1.13.1
pytest -v -s $TE_PATH/tests/*.py
pytest -v -s $TE_PATH/tests/test_transformerengine.py $TE_PATH/tests/test_fp8.py
NVTE_FLASH_ATTN=0 pytest -v -s $TE_PATH/tests/test_onnx_export.py
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

We don't test NVTE_FLASH_ATTN=1 I guess for ONNX?
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
 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
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Do we need L231-232 since we don't use layer_number and apply_query_key_layer_scaling in FA? I guess you're trying to keep the interface consistent between Unfused and Flash?
@ksivamanksivamanFeb 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.

Yes the second point is correct, just for consistency
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated


 if use_flash_attention:
 if checkpoint_core_attention:
 return self._checkpointed_attention_forward(query_layer, key_layer, value_layer)
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Did we forget to pass in the attention_func here? Same for L457.
@ksivamanksivamanFeb 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.

Yes there are some bugs in the selective activation checkpointing due to the MLM integration that doesn't actually run this path when I test for it, thus I wrongly assumed it works, fixing it now. Also will add standalone tests for this separately. Good catch :)
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
 custom_forward,
 False,
 self.get_rng_state_tracker,
 self.tp_group,
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Does self have a tp_group member?
@ksivamanksivamanFeb 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.

Same as above
@ksivaman
review comments and fixes to selective checkpointing
14a2beb
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 22, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
ksivaman
merged commit 79a9fe2
into
NVIDIA:mainFeb 22, 2023
@jeng1220jeng1220
mentioned this pull request
Feb 22, 2023
 Merged

@ksivamanksivaman
mentioned this pull request
Feb 24, 2023
 Closed

Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@ksivamanksivamanksivaman left review comments
@ptrendxptrendxptrendx approved these changes

Assignees


No one assigned


Labels


None yet


Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

4 participants

@cyanguwa@ksivaman@ptrendx@vasunvidia

, '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('^' + ".*" + '
flash-attn integration by cyanguwa · Pull Request #62 · NVIDIA/TransformerEngine · GitHub
Skip to content
flash-attn integration - #62
Merged
ksivaman merged 39 commits into
NVIDIA:mainfrom 
cyanguwa:flash_attn
Feb 22, 2023
Merged
flash-attn integration#62
ksivaman merged 39 commits into
NVIDIA:mainfrom 
cyanguwa:flash_attn
Conversation
@cyanguwa
@cyanguwacyanguwa
commented
Jan 31, 2023

Copy link


Copy Markdown

Collaborator
This PR is to provide support for flash attention in TE for self-attention calculations.
@ksivaman
ksivaman
marked this pull request as draft
January 31, 2023 00:53
@ksivaman
ksivaman
commented
Feb 4, 2023

Copy link


Copy Markdown

Member
@cyanguwa Could you sign off the commits? This guide explains how to do it.
@ksivaman
ksivaman
marked this pull request as ready for review
February 13, 2023 19:16
@ksivaman
ksivaman
self-requested a review
February 13, 2023 19:17
cyanguwaand others
added 13 commits
February 13, 2023 14:31
@cyanguwa
add flash attention to TransformerLayer
75ecfff
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Add docs for FP8 calibration (NVIDIA#61)
f428fe4
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ptrendx@cyanguwa
Fix the integer overflow in fused softmax (NVIDIA#60)
9c67222
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
prefix flash attn env var with NVTE_
faa7f5b
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Address steady memory increase and bloated checkpoints (NVIDIA#63)
ecef587
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
fix env var logic
1dccf39
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
fix flash attn env var logic again
38168dc
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
remove d2d copies (NVIDIA#64)
6ad62d3
* remove d2d copies
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* cleanup
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@vasunvidia@cyanguwa
Increase number of FP8 tensors per GEMM (NVIDIA#22)
f071530
* Increase number of FP8 tensors per GEMM
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Enable FP8 output tensor for fp8_gemm
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* [BERT FP8] Initial TE review comments
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Temporary fix for cuda graph non convergence
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Address review comments-2
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Review comments-3
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Cleanup
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Change for New API
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Remove unnecessary clone for D_scale, D_amax
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Avoid Roll for AMAX history size = 1
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Update onnx_te_gemm API
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Fix Lint errors
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
---------
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Bug fixes from PR 22 (NVIDIA#65)
d79c1b4
* Bug fixes from PR 22
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Add FP8 tests to ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* bundle unittests for ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
replace rearrange with transpose
be59386
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
QKV parameters unfused path fixes and optimization (NVIDIA#66)
faf9146
* Bug fixes from PR 22
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Add FP8 tests to ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Better QKV parameter fusion
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* small fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* keep original param for unfused case to retain externally set attrs
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* lint fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Fix ONNX exports
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* improve arg naming
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* No need to set data pointers
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* lint
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Assert memory loc in NoopCat
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Handle case of different memory in param and buffer
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* fix assert always true
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Reassign params memory to avoid more concats
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Fix gradients when using AMP (NVIDIA#70)
2416cc7
retain grad related attrs while casting
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
cyanguwaand others
added 8 commits
February 13, 2023 14:42
@cyanguwa
Merge branch 'main' into flash_attn
31cd706
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
fix pylint violations
d0f2eaf
fixed pyline violations such as trailing white spaces and too long lines Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
fix pylint violation on line 264 with R1719
304ef79
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
fix two more pylint violations
b3440c2
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
Merge branch 'NVIDIA:main' into flash_attn
cf28311
@ksivaman
DotProductAttention API
0b34943
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
Add docs for attention
39d4d71
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
fix assert always true
bd5160c
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 16, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
check for correct flash-attn version
4d8a6cb
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 16, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
review comments and fixes
28b571c
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
fix onnx and disable flash-attn export test
9b6b6c6
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
remove einops dependency
0551b43
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
cleanup internal API; rm duplication
cd6c90c
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
fix
9f82f42
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
only install TE wheel (exclude flash-attn to rm conflicts)
c9783b7
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 17, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
forgot to change install wheel path
ed64331
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
ptrendx
ptrendx
reviewed
Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated

 softmax_scale=1.0/self.norm_factor, causal=self.attn_causal_mask
 )

 # [b, sq, np, hn]
@ptrendxptrendxFeb 17, 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.

Wrong comment or wrong code?
@ksivamanksivamanFeb 17, 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.

The dimension comments are to indicate shape before not after
@ptrendxptrendxFeb 17, 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.

Well, then the code is wrong, since to get that you would need to do transpose, not view, right?
@ksivamanksivamanFeb 18, 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.

You are actually right, everything is wrong here; comment (input shape) and the view. I couldn't catch in my numerical tests since I'm running with mbs=1, sigh
ptrendx
ptrendx
reviewed
Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated

ptrendx
ptrendx
reviewed
Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
ptrendx
ptrendx
reviewed
Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated

@ksivaman
next round review comments
09687cc
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 17, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
fix flash_attn output
c39d043
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
ptrendx
ptrendx
approved these changes
Feb 19, 2023
View reviewed changes
@ptrendxptrendx
left a comment


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.

LGTM, thanks!
@ksivaman
fix QK layer scaling
1da83df
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
update docs
dc51968
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 22, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
ksivaman
commented
Feb 22, 2023

Copy link


Copy Markdown

Member
ci error is unrelated ONNX failure due to tolerance
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated

 ), 'FlashAttention currently only supports CUDA tensors.'
 assert (
 attention_mask is None
 ), 'FlashAttention currently does not support attention mask.'
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

FA calculates the causal attention mask inside itself so maybe here we meant 'FlashAttention doesn't support external attention mask'?
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadqa/L0_unittest/test.sh
pip install pytest==6.2.5 onnxruntime==1.13.1
pytest -v -s $TE_PATH/tests/*.py
pytest -v -s $TE_PATH/tests/test_transformerengine.py $TE_PATH/tests/test_fp8.py
NVTE_FLASH_ATTN=0 pytest -v -s $TE_PATH/tests/test_onnx_export.py
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

We don't test NVTE_FLASH_ATTN=1 I guess for ONNX?
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
 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
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Do we need L231-232 since we don't use layer_number and apply_query_key_layer_scaling in FA? I guess you're trying to keep the interface consistent between Unfused and Flash?
@ksivamanksivamanFeb 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.

Yes the second point is correct, just for consistency
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated


 if use_flash_attention:
 if checkpoint_core_attention:
 return self._checkpointed_attention_forward(query_layer, key_layer, value_layer)
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Did we forget to pass in the attention_func here? Same for L457.
@ksivamanksivamanFeb 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.

Yes there are some bugs in the selective activation checkpointing due to the MLM integration that doesn't actually run this path when I test for it, thus I wrongly assumed it works, fixing it now. Also will add standalone tests for this separately. Good catch :)
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
 custom_forward,
 False,
 self.get_rng_state_tracker,
 self.tp_group,
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Does self have a tp_group member?
@ksivamanksivamanFeb 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.

Same as above
@ksivaman
review comments and fixes to selective checkpointing
14a2beb
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 22, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
ksivaman
merged commit 79a9fe2
into
NVIDIA:mainFeb 22, 2023
@jeng1220jeng1220
mentioned this pull request
Feb 22, 2023
 Merged

@ksivamanksivaman
mentioned this pull request
Feb 24, 2023
 Closed

Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@ksivamanksivamanksivaman left review comments
@ptrendxptrendxptrendx approved these changes

Assignees


No one assigned


Labels


None yet


Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

4 participants

@cyanguwa@ksivaman@ptrendx@vasunvidia

, '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('^' + ".*" + '
flash-attn integration by cyanguwa · Pull Request #62 · NVIDIA/TransformerEngine · GitHub
Skip to content
flash-attn integration - #62
Merged
ksivaman merged 39 commits into
NVIDIA:mainfrom 
cyanguwa:flash_attn
Feb 22, 2023
Merged
flash-attn integration#62
ksivaman merged 39 commits into
NVIDIA:mainfrom 
cyanguwa:flash_attn
Conversation
@cyanguwa
@cyanguwacyanguwa
commented
Jan 31, 2023

Copy link


Copy Markdown

Collaborator
This PR is to provide support for flash attention in TE for self-attention calculations.
@ksivaman
ksivaman
marked this pull request as draft
January 31, 2023 00:53
@ksivaman
ksivaman
commented
Feb 4, 2023

Copy link


Copy Markdown

Member
@cyanguwa Could you sign off the commits? This guide explains how to do it.
@ksivaman
ksivaman
marked this pull request as ready for review
February 13, 2023 19:16
@ksivaman
ksivaman
self-requested a review
February 13, 2023 19:17
cyanguwaand others
added 13 commits
February 13, 2023 14:31
@cyanguwa
add flash attention to TransformerLayer
75ecfff
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Add docs for FP8 calibration (NVIDIA#61)
f428fe4
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ptrendx@cyanguwa
Fix the integer overflow in fused softmax (NVIDIA#60)
9c67222
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
prefix flash attn env var with NVTE_
faa7f5b
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Address steady memory increase and bloated checkpoints (NVIDIA#63)
ecef587
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
fix env var logic
1dccf39
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
fix flash attn env var logic again
38168dc
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
remove d2d copies (NVIDIA#64)
6ad62d3
* remove d2d copies
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* cleanup
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@vasunvidia@cyanguwa
Increase number of FP8 tensors per GEMM (NVIDIA#22)
f071530
* Increase number of FP8 tensors per GEMM
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Enable FP8 output tensor for fp8_gemm
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* [BERT FP8] Initial TE review comments
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Temporary fix for cuda graph non convergence
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Address review comments-2
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Review comments-3
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Cleanup
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Change for New API
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Remove unnecessary clone for D_scale, D_amax
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Avoid Roll for AMAX history size = 1
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Update onnx_te_gemm API
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Fix Lint errors
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
---------
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Bug fixes from PR 22 (NVIDIA#65)
d79c1b4
* Bug fixes from PR 22
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Add FP8 tests to ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* bundle unittests for ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
replace rearrange with transpose
be59386
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
QKV parameters unfused path fixes and optimization (NVIDIA#66)
faf9146
* Bug fixes from PR 22
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Add FP8 tests to ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Better QKV parameter fusion
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* small fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* keep original param for unfused case to retain externally set attrs
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* lint fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Fix ONNX exports
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* improve arg naming
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* No need to set data pointers
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* lint
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Assert memory loc in NoopCat
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Handle case of different memory in param and buffer
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* fix assert always true
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Reassign params memory to avoid more concats
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Fix gradients when using AMP (NVIDIA#70)
2416cc7
retain grad related attrs while casting
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
cyanguwaand others
added 8 commits
February 13, 2023 14:42
@cyanguwa
Merge branch 'main' into flash_attn
31cd706
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
fix pylint violations
d0f2eaf
fixed pyline violations such as trailing white spaces and too long lines Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
fix pylint violation on line 264 with R1719
304ef79
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
fix two more pylint violations
b3440c2
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
Merge branch 'NVIDIA:main' into flash_attn
cf28311
@ksivaman
DotProductAttention API
0b34943
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
Add docs for attention
39d4d71
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
fix assert always true
bd5160c
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 16, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
check for correct flash-attn version
4d8a6cb
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 16, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
review comments and fixes
28b571c
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
fix onnx and disable flash-attn export test
9b6b6c6
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
remove einops dependency
0551b43
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
cleanup internal API; rm duplication
cd6c90c
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
fix
9f82f42
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
only install TE wheel (exclude flash-attn to rm conflicts)
c9783b7
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 17, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
forgot to change install wheel path
ed64331
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
ptrendx
ptrendx
reviewed
Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated

 softmax_scale=1.0/self.norm_factor, causal=self.attn_causal_mask
 )

 # [b, sq, np, hn]
@ptrendxptrendxFeb 17, 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.

Wrong comment or wrong code?
@ksivamanksivamanFeb 17, 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.

The dimension comments are to indicate shape before not after
@ptrendxptrendxFeb 17, 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.

Well, then the code is wrong, since to get that you would need to do transpose, not view, right?
@ksivamanksivamanFeb 18, 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.

You are actually right, everything is wrong here; comment (input shape) and the view. I couldn't catch in my numerical tests since I'm running with mbs=1, sigh
ptrendx
ptrendx
reviewed
Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated

ptrendx
ptrendx
reviewed
Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
ptrendx
ptrendx
reviewed
Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated

@ksivaman
next round review comments
09687cc
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 17, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
fix flash_attn output
c39d043
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
ptrendx
ptrendx
approved these changes
Feb 19, 2023
View reviewed changes
@ptrendxptrendx
left a comment


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.

LGTM, thanks!
@ksivaman
fix QK layer scaling
1da83df
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
update docs
dc51968
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 22, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
ksivaman
commented
Feb 22, 2023

Copy link


Copy Markdown

Member
ci error is unrelated ONNX failure due to tolerance
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated

 ), 'FlashAttention currently only supports CUDA tensors.'
 assert (
 attention_mask is None
 ), 'FlashAttention currently does not support attention mask.'
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

FA calculates the causal attention mask inside itself so maybe here we meant 'FlashAttention doesn't support external attention mask'?
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadqa/L0_unittest/test.sh
pip install pytest==6.2.5 onnxruntime==1.13.1
pytest -v -s $TE_PATH/tests/*.py
pytest -v -s $TE_PATH/tests/test_transformerengine.py $TE_PATH/tests/test_fp8.py
NVTE_FLASH_ATTN=0 pytest -v -s $TE_PATH/tests/test_onnx_export.py
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

We don't test NVTE_FLASH_ATTN=1 I guess for ONNX?
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
 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
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Do we need L231-232 since we don't use layer_number and apply_query_key_layer_scaling in FA? I guess you're trying to keep the interface consistent between Unfused and Flash?
@ksivamanksivamanFeb 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.

Yes the second point is correct, just for consistency
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated


 if use_flash_attention:
 if checkpoint_core_attention:
 return self._checkpointed_attention_forward(query_layer, key_layer, value_layer)
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Did we forget to pass in the attention_func here? Same for L457.
@ksivamanksivamanFeb 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.

Yes there are some bugs in the selective activation checkpointing due to the MLM integration that doesn't actually run this path when I test for it, thus I wrongly assumed it works, fixing it now. Also will add standalone tests for this separately. Good catch :)
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
 custom_forward,
 False,
 self.get_rng_state_tracker,
 self.tp_group,
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Does self have a tp_group member?
@ksivamanksivamanFeb 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.

Same as above
@ksivaman
review comments and fixes to selective checkpointing
14a2beb
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 22, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
ksivaman
merged commit 79a9fe2
into
NVIDIA:mainFeb 22, 2023
@jeng1220jeng1220
mentioned this pull request
Feb 22, 2023
 Merged

@ksivamanksivaman
mentioned this pull request
Feb 24, 2023
 Closed

Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@ksivamanksivamanksivaman left review comments
@ptrendxptrendxptrendx approved these changes

Assignees


No one assigned


Labels


None yet


Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

4 participants

@cyanguwa@ksivaman@ptrendx@vasunvidia

, '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); }
})();
})();
flash-attn integration by cyanguwa · Pull Request #62 · NVIDIA/TransformerEngine · GitHub
Skip to content
flash-attn integration - #62
Merged
ksivaman merged 39 commits into
NVIDIA:mainfrom 
cyanguwa:flash_attn
Feb 22, 2023
Merged
flash-attn integration#62
ksivaman merged 39 commits into
NVIDIA:mainfrom 
cyanguwa:flash_attn
Conversation
@cyanguwa
@cyanguwacyanguwa
commented
Jan 31, 2023

Copy link


Copy Markdown

Collaborator
This PR is to provide support for flash attention in TE for self-attention calculations.
@ksivaman
ksivaman
marked this pull request as draft
January 31, 2023 00:53
@ksivaman
ksivaman
commented
Feb 4, 2023

Copy link


Copy Markdown

Member
@cyanguwa Could you sign off the commits? This guide explains how to do it.
@ksivaman
ksivaman
marked this pull request as ready for review
February 13, 2023 19:16
@ksivaman
ksivaman
self-requested a review
February 13, 2023 19:17
cyanguwaand others
added 13 commits
February 13, 2023 14:31
@cyanguwa
add flash attention to TransformerLayer
75ecfff
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Add docs for FP8 calibration (NVIDIA#61)
f428fe4
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ptrendx@cyanguwa
Fix the integer overflow in fused softmax (NVIDIA#60)
9c67222
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
prefix flash attn env var with NVTE_
faa7f5b
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Address steady memory increase and bloated checkpoints (NVIDIA#63)
ecef587
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
fix env var logic
1dccf39
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
fix flash attn env var logic again
38168dc
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
remove d2d copies (NVIDIA#64)
6ad62d3
* remove d2d copies
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* cleanup
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@vasunvidia@cyanguwa
Increase number of FP8 tensors per GEMM (NVIDIA#22)
f071530
* Increase number of FP8 tensors per GEMM
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Enable FP8 output tensor for fp8_gemm
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* [BERT FP8] Initial TE review comments
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Temporary fix for cuda graph non convergence
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Address review comments-2
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Review comments-3
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Cleanup
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Change for New API
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Remove unnecessary clone for D_scale, D_amax
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Avoid Roll for AMAX history size = 1
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Update onnx_te_gemm API
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
* Fix Lint errors
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
---------
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Bug fixes from PR 22 (NVIDIA#65)
d79c1b4
* Bug fixes from PR 22
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Add FP8 tests to ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* bundle unittests for ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@cyanguwa
replace rearrange with transpose
be59386
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
QKV parameters unfused path fixes and optimization (NVIDIA#66)
faf9146
* Bug fixes from PR 22
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Add FP8 tests to ci
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Better QKV parameter fusion
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* small fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* keep original param for unfused case to retain externally set attrs
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* lint fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Fix ONNX exports
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* improve arg naming
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* No need to set data pointers
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* lint
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Assert memory loc in NoopCat
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Handle case of different memory in param and buffer
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* fix assert always true
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Reassign params memory to avoid more concats
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@ksivaman@cyanguwa
Fix gradients when using AMP (NVIDIA#70)
2416cc7
retain grad related attrs while casting
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
cyanguwaand others
added 8 commits
February 13, 2023 14:42
@cyanguwa
Merge branch 'main' into flash_attn
31cd706
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
fix pylint violations
d0f2eaf
fixed pyline violations such as trailing white spaces and too long lines Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
fix pylint violation on line 264 with R1719
304ef79
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
fix two more pylint violations
b3440c2
Signed-off-by: cyanguwa <cyang.uwa@gmail.com>
@cyanguwa
Merge branch 'NVIDIA:main' into flash_attn
cf28311
@ksivaman
DotProductAttention API
0b34943
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
Add docs for attention
39d4d71
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
fix assert always true
bd5160c
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 16, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
check for correct flash-attn version
4d8a6cb
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 16, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
review comments and fixes
28b571c
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
fix onnx and disable flash-attn export test
9b6b6c6
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
remove einops dependency
0551b43
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
cleanup internal API; rm duplication
cd6c90c
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
fix
9f82f42
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
only install TE wheel (exclude flash-attn to rm conflicts)
c9783b7
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 17, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
forgot to change install wheel path
ed64331
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
ptrendx
ptrendx
reviewed
Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated

 softmax_scale=1.0/self.norm_factor, causal=self.attn_causal_mask
 )

 # [b, sq, np, hn]
@ptrendxptrendxFeb 17, 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.

Wrong comment or wrong code?
@ksivamanksivamanFeb 17, 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.

The dimension comments are to indicate shape before not after
@ptrendxptrendxFeb 17, 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.

Well, then the code is wrong, since to get that you would need to do transpose, not view, right?
@ksivamanksivamanFeb 18, 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.

You are actually right, everything is wrong here; comment (input shape) and the view. I couldn't catch in my numerical tests since I'm running with mbs=1, sigh
ptrendx
ptrendx
reviewed
Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated

ptrendx
ptrendx
reviewed
Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
ptrendx
ptrendx
reviewed
Feb 17, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated

@ksivaman
next round review comments
09687cc
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 17, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
fix flash_attn output
c39d043
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
ptrendx
ptrendx
approved these changes
Feb 19, 2023
View reviewed changes
@ptrendxptrendx
left a comment


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.

LGTM, thanks!
@ksivaman
fix QK layer scaling
1da83df
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
update docs
dc51968
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 22, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
ksivaman
commented
Feb 22, 2023

Copy link


Copy Markdown

Member
ci error is unrelated ONNX failure due to tolerance
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated

 ), 'FlashAttention currently only supports CUDA tensors.'
 assert (
 attention_mask is None
 ), 'FlashAttention currently does not support attention mask.'
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

FA calculates the causal attention mask inside itself so maybe here we meant 'FlashAttention doesn't support external attention mask'?
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadqa/L0_unittest/test.sh
pip install pytest==6.2.5 onnxruntime==1.13.1
pytest -v -s $TE_PATH/tests/*.py
pytest -v -s $TE_PATH/tests/test_transformerengine.py $TE_PATH/tests/test_fp8.py
NVTE_FLASH_ATTN=0 pytest -v -s $TE_PATH/tests/test_onnx_export.py
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

We don't test NVTE_FLASH_ATTN=1 I guess for ONNX?
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
 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
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Do we need L231-232 since we don't use layer_number and apply_query_key_layer_scaling in FA? I guess you're trying to keep the interface consistent between Unfused and Flash?
@ksivamanksivamanFeb 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.

Yes the second point is correct, just for consistency
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
Outdated


 if use_flash_attention:
 if checkpoint_core_attention:
 return self._checkpointed_attention_forward(query_layer, key_layer, value_layer)
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Did we forget to pass in the attention_func here? Same for L457.
@ksivamanksivamanFeb 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.

Yes there are some bugs in the selective activation checkpointing due to the MLM integration that doesn't actually run this path when I test for it, thus I wrongly assumed it works, fixing it now. Also will add standalone tests for this separately. Good catch :)
cyanguwa
cyanguwa
commented
Feb 22, 2023
View reviewed changes
Comment threadtransformer_engine/pytorch/transformer.py
 custom_forward,
 False,
 self.get_rng_state_tracker,
 self.tp_group,
@cyanguwacyanguwaFeb 22, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Does self have a tp_group member?
@ksivamanksivamanFeb 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.

Same as above
@ksivaman
review comments and fixes to selective checkpointing
14a2beb
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman
ksivaman
commented
Feb 22, 2023

Copy link


Copy Markdown

Member
/te-ci
@ksivaman
ksivaman
merged commit 79a9fe2
into
NVIDIA:mainFeb 22, 2023
@jeng1220jeng1220
mentioned this pull request
Feb 22, 2023
 Merged

@ksivamanksivaman
mentioned this pull request
Feb 24, 2023
 Closed

Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@ksivamanksivamanksivaman left review comments
@ptrendxptrendxptrendx approved these changes

Assignees


No one assigned


Labels


None yet


Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

4 participants

@cyanguwa@ksivaman@ptrendx@vasunvidia