') + ')', '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('^' + ".*" + ', '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" + ', '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('^' + ".*" + ', '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); } })(); })(); [JAX] Add self_attn_mask_type and replace attn_type by zlsh80826 · Pull Request #273 · NVIDIA/TransformerEngine · GitHub
Skip to content

[JAX] Add self_attn_mask_type and replace attn_type - #273

Merged
ksivaman merged 5 commits into
NVIDIA:mainfrom
zlsh80826:rewang/config-self-attn-mask-type
Jun 20, 2023
Merged

[JAX] Add self_attn_mask_type and replace attn_type#273
ksivaman merged 5 commits into
NVIDIA:mainfrom
zlsh80826:rewang/config-self-attn-mask-type

Conversation

@zlsh80826

@zlsh80826zlsh80826 commented Jun 12, 2023•
edited
Loading

Copy link
Copy Markdown
Collaborator

This PR does the following two changes

  1. Configuration of self_attn_mask_type in TransformerLayer:
    The current implementation of TransformerLayer has a hardcoded self-attention mask type, which follows the behavior of T5-like models. It uses 'padding' for the encoder and 'causal' for the decoder. However, this behavior may not be applicable for all transformer variant models. To address this limitation, this PR introduces the self_attn_mask_type as an argument in the TransformerLayer. We plan to make the default value for self_attn_mask_type be 'causal' (aligned with pyTorch) in the next two releases (0.11+). Prior to that, the default behavior will remain the same as the old version, using 'padding' for the encoder and 'causal' for the decoder. A FutureWarning will be generated to notify users about the upcoming change.

  2. Deprecation of attn_type and addition of attn_mask_type in MultiHeadAttention:
    The current naming of attn_type is ambiguous. To address this, this PR replaces attn_type with attn_mask_type in the MultiHeadAttention module. It is worth noting that the impact on end users should be minimal or even non-existent, as we inadvertently did not export the internal AttentionType in previous versions. However, we will still keep attn_type in the upcoming release (0.10) while ignoring its usage, with the plan to completely remove it in the subsequent two releases (0.11+).

All of the above changes are also added into the documentation and also the runtime warnings.

@zlsh80826
zlsh80826force-pushed the rewang/config-self-attn-mask-type branch from 954a782 to e1e41a6CompareJune 12, 2023 09:49
@zlsh80826

zlsh80826 commented Jun 12, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

@zlsh80826
zlsh80826 marked this pull request as ready for review June 12, 2023 12:34
@zlsh80826zlsh80826 added documentation Improvements or additions to documentation enhancement New feature or request labels Jun 12, 2023
@zlsh80826
Add self_attn_mask_type and replace attn_type
2176178
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Refine the keyword style for the better readability
e801beb
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
Replace attn_type with attn_mask_type in praxis transformer
e1404b6
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826
zlsh80826force-pushed the rewang/config-self-attn-mask-type branch from bd1547c to e1404b6CompareJune 13, 2023 05:12
@zlsh80826

zlsh80826 commented Jun 13, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

ksivaman
ksivaman reviewed Jun 13, 2023
View reviewed changes
Comment threadtransformer_engine/jax/flax/transformer.py Outdated
ksivaman
ksivaman reviewed Jun 13, 2023
View reviewed changes
Comment threadtransformer_engine/jax/flax/transformer.py Outdated
ksivaman
ksivaman reviewed Jun 13, 2023
View reviewed changes
Comment threadtransformer_engine/jax/flax/transformer.py Outdated
ksivaman
ksivaman reviewed Jun 13, 2023
View reviewed changes
Comment threadtransformer_engine/jax/flax/transformer.py Outdated
ksivaman
ksivaman approved these changes Jun 13, 2023
View reviewed changes

@ksivamanksivaman 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 other than small suggestions

@zlsh80826
Fix typos
20ab5c0
Signed-off-by: Reese Wang <rewang@nvidia.com>
@zlsh80826

zlsh80826 commented Jun 13, 2023

Copy link
Copy Markdown
CollaboratorAuthor

Thank @ksivaman for catching the typos!

@zlsh80826

zlsh80826 commented Jun 13, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

mingxu1067
mingxu1067 approved these changes Jun 14, 2023
View reviewed changes

@mingxu1067mingxu1067 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

LGTM

@zlsh80826

zlsh80826 commented Jun 20, 2023•
edited
Loading

Copy link
Copy Markdown
CollaboratorAuthor

Hi @ksivaman, I saw you freezed the code for branch 0.10, could this PR also be cherry-picked into branch 0.10? This PR is a fix for the PAXML/GPT.

@ksivaman

ksivaman commented Jun 20, 2023

Copy link
Copy Markdown
Member

/te-ci

@ksivaman
ksivaman merged commit 16208b3 into NVIDIA:mainJun 20, 2023
@ksivaman

ksivaman commented Jun 20, 2023

Copy link
Copy Markdown
Member

Pipeline 8688979

ksivaman added a commit that referenced this pull request Jun 22, 2023
@zlsh80826@ksivaman
[JAX] Add self_attn_mask_type and replace attn_type (#273)
84a4a75
* Add self_attn_mask_type and replace attn_type
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Refine the keyword style for the better readability
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Replace attn_type with attn_mask_type in praxis transformer
Signed-off-by: Reese Wang <rewang@nvidia.com>
* Fix typos
Signed-off-by: Reese Wang <rewang@nvidia.com>
---------
Signed-off-by: Reese Wang <rewang@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@zlsh80826zlsh80826 mentioned this pull request Jul 16, 2023
Merged
@hungryGeek16hungryGeek16 mentioned this pull request May 31, 2026
Open
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ksivamanksivamanksivaman approved these changes
@timmoon10timmoon10Awaiting requested review from timmoon10
@ptrendxptrendxAwaiting requested review from ptrendx
+1 more reviewer
@mingxu1067mingxu1067mingxu1067 approved these changes
Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

0.10.0documentationImprovements or additions to documentationenhancementNew feature or request

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@zlsh80826@ksivaman@mingxu1067