Skip to content

The Implementation of Praxis's Modules - #158

Merged
timmoon10 merged 7 commits into
NVIDIA:mainfrom
mingxu1067:mingh/praxis_modules_for_review
May 16, 2023
Merged

The Implementation of Praxis's Modules#158
timmoon10 merged 7 commits into
NVIDIA:mainfrom
mingxu1067:mingh/praxis_modules_for_review

Conversation

@mingxu1067

@mingxu1067mingxu1067 commented Apr 19, 2023

Copy link
Copy Markdown
Collaborator

No description provided.

@mingxu1067

mingxu1067 commented Apr 19, 2023

Copy link
Copy Markdown
CollaboratorAuthor

@nouiz and @timmoon10 , this is our initial implementation to support Praxis, and this is built on top of #152 .

@mingxu1067

mingxu1067 commented Apr 19, 2023

Copy link
Copy Markdown
CollaboratorAuthor

@jeng1220 for viz.

@ksivaman
ksivaman marked this pull request as draft April 19, 2023 03:56
@ksivamanksivaman changed the title [WIP, DO NOT MERGE] The Implementation of Praxis's Modules for Pre-reviewThe Implementation of Praxis's Modules for Pre-reviewApr 19, 2023
nouiz
nouiz reviewed Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/jax/praxis/module.py Outdated
@mingxu1067
mingxu1067force-pushed the mingh/praxis_modules_for_review branch from db8114d to 910948bCompareApril 21, 2023 03:22
@mingxu1067

mingxu1067 commented Apr 21, 2023

Copy link
Copy Markdown
CollaboratorAuthor

Update:

  1. Rebased on main.
  2. Added Zero-Centered-Gamma supports to LayerNorm.

nouiz
nouiz reviewed Apr 21, 2023
View reviewed changes
Comment threadtransformer_engine/jax/praxis/transformer.py Outdated
@mingxu1067
mingxu1067force-pushed the mingh/praxis_modules_for_review branch from 910948b to c4fb418CompareApril 24, 2023 05:50
@mingxu1067

mingxu1067 commented Apr 24, 2023

Copy link
Copy Markdown
CollaboratorAuthor

Update:

  1. Rebased on [JAX] Adjust Module Structure. #169 .
  2. Added Zero-Centered-Gamma supports to all LN-related modules.

timmoon10
timmoon10 reviewed Apr 24, 2023
View reviewed changes

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

This PR looks reasonable to me, pending the merge of #169 and fixing test failures.

Comment threadtests/jax/test_praxis_layers.py Outdated

@timmoon10timmoon10Apr 24, 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.

These tests are fine, but they could be made more challenging by comparing the Praxis layers against the non-TE implementations in https://github.com/NVIDIA/TransformerEngine/blob/main/tests/jax/utils.py. After all, the Praxis layers are mostly generated from the Flax layers anyways. That said, this is not high priority.

@mingxu1067mingxu1067May 2, 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.

Agree, put into TODO.

nouiz
nouiz reviewed May 1, 2023
View reviewed changes
Comment threadtransformer_engine/jax/praxis/module.py Outdated

@nouiznouizMay 1, 2023

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

The class name is different in praxis and flax.
Any reason for that?

@mingxu1067mingxu1067May 2, 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.

FusedSoftmax is more sutiable to its behavior. We had a thought to change the name in Flax. But haven't have a conclusion yet.

@mingxu1067
mingxu1067force-pushed the mingh/praxis_modules_for_review branch from c4fb418 to dc4dfc1CompareMay 2, 2023 03:22
@mingxu1067

mingxu1067 commented May 2, 2023

Copy link
Copy Markdown
CollaboratorAuthor

Update:

  1. Adapt to the latest Praxis.

@timmoon10
timmoon10 self-requested a review May 3, 2023 23:12
timmoon10
timmoon10 approved these changes May 4, 2023
View reviewed changes

@timmoon10timmoon10 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 once we fix the test failures.

@nouiz

nouiz commented May 8, 2023•
edited
Loading

Copy link
Copy Markdown
Collaborator

@mingxu1067 as we only need to fix the test failure, can you remove the draft status?

@mingxu1067

mingxu1067 commented May 9, 2023•
edited
Loading

Copy link
Copy Markdown
CollaboratorAuthor

@nouiz , I will rebase this on top of #196 then mark as ready.

@mingxu1067
mingxu1067force-pushed the mingh/praxis_modules_for_review branch from dc4dfc1 to 116ce7aCompareMay 10, 2023 02:03
@mingxu1067
mingxu1067 marked this pull request as ready for review May 10, 2023 02:05
@mingxu1067

mingxu1067 commented May 10, 2023

Copy link
Copy Markdown
CollaboratorAuthor

Rebased on main.

timmoon10
timmoon10 requested changes May 12, 2023
View reviewed changes
Comment threadtransformer_engine/jax/__init__.py Outdated
@mingxu1067
Adding JAX/Praxis modules and dependencies.
b75b0b9
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Adding UTs to JAX/Praxis modules.
23410e2
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Remove praxis as a dependency due to not strictly needed
08d09ab
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Repalce is_fp8_supported to is_fp8_available
11a8e66
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Make Praxis as an optional dependency.
4823c01
1. Removed 'from . import praxis' in __init__.py.
1.1 Noted, keep 'from . import flax' for deprecated warning.
2. Changed te.flax to te_flax in examples and README.rst.
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
mingxu1067force-pushed the mingh/praxis_modules_for_review branch from cadc4b0 to 4823c01CompareMay 15, 2023 02:53
@mingxu1067

mingxu1067 commented May 15, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

@mingxu1067mingxu1067 changed the title The Implementation of Praxis's Modules for Pre-reviewThe Implementation of Praxis's ModulesMay 15, 2023
@mingxu1067
Adding a workaround to FP8 training on Praxis.
c72007e
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067

mingxu1067 commented May 15, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

@mingxu1067

mingxu1067 commented May 15, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

timmoon10
timmoon10 approved these changes May 15, 2023
View reviewed changes

@timmoon10timmoon10 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

@mingxu1067

mingxu1067 commented May 16, 2023

Copy link
Copy Markdown
CollaboratorAuthor

CI is passed, could we merge this? Thanks.

@timmoon10
timmoon10 merged commit b20c053 into NVIDIA:mainMay 16, 2023
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@timmoon10timmoon10timmoon10 approved these changes
+1 more reviewer
@nouiznouiznouiz left review comments
Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@mingxu1067@nouiz@timmoon10
, '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" + '
The Implementation of Praxis's Modules by mingxu1067 · Pull Request #158 · NVIDIA/TransformerEngine · GitHub
Skip to content
The Implementation of Praxis's Modules - #158
Merged
timmoon10 merged 7 commits into
NVIDIA:mainfrom 
mingxu1067:mingh/praxis_modules_for_review
May 16, 2023
Merged
The Implementation of Praxis's Modules#158
timmoon10 merged 7 commits into
NVIDIA:mainfrom 
mingxu1067:mingh/praxis_modules_for_review
Conversation
@mingxu1067
@mingxu1067mingxu1067
commented
Apr 19, 2023

Copy link


Copy Markdown

Collaborator
 No description provided. 
@mingxu1067
mingxu1067
commented
Apr 19, 2023

Copy link


Copy Markdown

CollaboratorAuthor
@nouiz and @timmoon10 , this is our initial implementation to support Praxis, and this is built on top of #152 .
@mingxu1067
mingxu1067
commented
Apr 19, 2023

Copy link


Copy Markdown

CollaboratorAuthor
@jeng1220 for viz.
@ksivaman
ksivaman
marked this pull request as draft
April 19, 2023 03:56
@ksivamanksivaman
changed the title
[WIP, DO NOT MERGE] The Implementation of Praxis's Modules for Pre-reviewThe Implementation of Praxis's Modules for Pre-reviewApr 19, 2023
nouiz
nouiz
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/jax/praxis/module.py
Outdated

@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
db8114d to
910948bCompareApril 21, 2023 03:22
@mingxu1067
mingxu1067
commented
Apr 21, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Update:
  1. Rebased on main.
  2. Added Zero-Centered-Gamma supports to LayerNorm.
nouiz
nouiz
reviewed
Apr 21, 2023
View reviewed changes
Comment threadtransformer_engine/jax/praxis/transformer.py
Outdated

@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
910948b to
c4fb418CompareApril 24, 2023 05:50
@mingxu1067
mingxu1067
commented
Apr 24, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Update:
  1. Rebased on [JAX] Adjust Module Structure. #169 .
  2. Added Zero-Centered-Gamma supports to all LN-related modules.
timmoon10
timmoon10
reviewed
Apr 24, 2023
View reviewed changes
@timmoon10timmoon10
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.

This PR looks reasonable to me, pending the merge of #169 and fixing test failures.
Comment threadtests/jax/test_praxis_layers.py
Outdated

@timmoon10timmoon10Apr 24, 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.

These tests are fine, but they could be made more challenging by comparing the Praxis layers against the non-TE implementations in https://github.com/NVIDIA/TransformerEngine/blob/main/tests/jax/utils.py. After all, the Praxis layers are mostly generated from the Flax layers anyways. That said, this is not high priority.
@mingxu1067mingxu1067May 2, 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.

Agree, put into TODO.
nouiz
nouiz
reviewed
May 1, 2023
View reviewed changes
Comment threadtransformer_engine/jax/praxis/module.py
Outdated

@nouiznouizMay 1, 2023

Copy link


Copy Markdown

Collaborator
Choose a reason for hiding this comment

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

The class name is different in praxis and flax.

Any reason for that?
@mingxu1067mingxu1067May 2, 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.

FusedSoftmax is more sutiable to its behavior. We had a thought to change the name in Flax. But haven't have a conclusion yet.
@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
c4fb418 to
dc4dfc1CompareMay 2, 2023 03:22
@mingxu1067
mingxu1067
commented
May 2, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Update:
  1. Adapt to the latest Praxis.
@timmoon10
timmoon10
self-requested a review
May 3, 2023 23:12
timmoon10
timmoon10
approved these changes
May 4, 2023
View reviewed changes
@timmoon10timmoon10
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 once we fix the test failures.
@nouiz
nouiz
commented
May 8, 2023•

edited

Loading

Copy link


Copy Markdown

Collaborator
@mingxu1067 as we only need to fix the test failure, can you remove the draft status?
@mingxu1067
mingxu1067
commented
May 9, 2023•

edited

Loading

Copy link


Copy Markdown

CollaboratorAuthor
@nouiz , I will rebase this on top of #196 then mark as ready.
@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
dc4dfc1 to
116ce7aCompareMay 10, 2023 02:03
@mingxu1067
mingxu1067
marked this pull request as ready for review
May 10, 2023 02:05
@mingxu1067
mingxu1067
commented
May 10, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Rebased on main.
timmoon10
timmoon10
requested changes
May 12, 2023
View reviewed changes
Comment threadtransformer_engine/jax/__init__.py
Outdated

@mingxu1067
Adding JAX/Praxis modules and dependencies.
b75b0b9
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Adding UTs to JAX/Praxis modules.
23410e2
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Remove praxis as a dependency due to not strictly needed
08d09ab
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Repalce is_fp8_supported to is_fp8_available
11a8e66
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Make Praxis as an optional dependency.
4823c01
1. Removed 'from . import praxis' in __init__.py.
1.1 Noted, keep 'from . import flax' for deprecated warning.
2. Changed te.flax to te_flax in examples and README.rst.
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
cadc4b0 to
4823c01CompareMay 15, 2023 02:53
@mingxu1067
mingxu1067
commented
May 15, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@mingxu1067mingxu1067
changed the title
The Implementation of Praxis's Modules for Pre-reviewThe Implementation of Praxis's ModulesMay 15, 2023
@mingxu1067
Adding a workaround to FP8 training on Praxis.
c72007e
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
mingxu1067
commented
May 15, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@mingxu1067
mingxu1067
commented
May 15, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
timmoon10
timmoon10
approved these changes
May 15, 2023
View reviewed changes
@timmoon10timmoon10
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
@mingxu1067
mingxu1067
commented
May 16, 2023

Copy link


Copy Markdown

CollaboratorAuthor
CI is passed, could we merge this? Thanks.
@timmoon10
timmoon10
merged commit b20c053
into
NVIDIA:mainMay 16, 2023
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@timmoon10timmoon10timmoon10 approved these changes

+1 more reviewer

@nouiznouiznouiz left review comments
Reviewers whose approvals may not affect merge requirements

Assignees


No one assigned


Labels


None yet


Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

3 participants

@mingxu1067@nouiz@timmoon10

, '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('^' + ".*" + '
The Implementation of Praxis's Modules by mingxu1067 · Pull Request #158 · NVIDIA/TransformerEngine · GitHub
Skip to content
The Implementation of Praxis's Modules - #158
Merged
timmoon10 merged 7 commits into
NVIDIA:mainfrom 
mingxu1067:mingh/praxis_modules_for_review
May 16, 2023
Merged
The Implementation of Praxis's Modules#158
timmoon10 merged 7 commits into
NVIDIA:mainfrom 
mingxu1067:mingh/praxis_modules_for_review
Conversation
@mingxu1067
@mingxu1067mingxu1067
commented
Apr 19, 2023

Copy link


Copy Markdown

Collaborator
 No description provided. 
@mingxu1067
mingxu1067
commented
Apr 19, 2023

Copy link


Copy Markdown

CollaboratorAuthor
@nouiz and @timmoon10 , this is our initial implementation to support Praxis, and this is built on top of #152 .
@mingxu1067
mingxu1067
commented
Apr 19, 2023

Copy link


Copy Markdown

CollaboratorAuthor
@jeng1220 for viz.
@ksivaman
ksivaman
marked this pull request as draft
April 19, 2023 03:56
@ksivamanksivaman
changed the title
[WIP, DO NOT MERGE] The Implementation of Praxis's Modules for Pre-reviewThe Implementation of Praxis's Modules for Pre-reviewApr 19, 2023
nouiz
nouiz
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/jax/praxis/module.py
Outdated

@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
db8114d to
910948bCompareApril 21, 2023 03:22
@mingxu1067
mingxu1067
commented
Apr 21, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Update:
  1. Rebased on main.
  2. Added Zero-Centered-Gamma supports to LayerNorm.
nouiz
nouiz
reviewed
Apr 21, 2023
View reviewed changes
Comment threadtransformer_engine/jax/praxis/transformer.py
Outdated

@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
910948b to
c4fb418CompareApril 24, 2023 05:50
@mingxu1067
mingxu1067
commented
Apr 24, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Update:
  1. Rebased on [JAX] Adjust Module Structure. #169 .
  2. Added Zero-Centered-Gamma supports to all LN-related modules.
timmoon10
timmoon10
reviewed
Apr 24, 2023
View reviewed changes
@timmoon10timmoon10
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.

This PR looks reasonable to me, pending the merge of #169 and fixing test failures.
Comment threadtests/jax/test_praxis_layers.py
Outdated

@timmoon10timmoon10Apr 24, 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.

These tests are fine, but they could be made more challenging by comparing the Praxis layers against the non-TE implementations in https://github.com/NVIDIA/TransformerEngine/blob/main/tests/jax/utils.py. After all, the Praxis layers are mostly generated from the Flax layers anyways. That said, this is not high priority.
@mingxu1067mingxu1067May 2, 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.

Agree, put into TODO.
nouiz
nouiz
reviewed
May 1, 2023
View reviewed changes
Comment threadtransformer_engine/jax/praxis/module.py
Outdated

@nouiznouizMay 1, 2023

Copy link


Copy Markdown

Collaborator
Choose a reason for hiding this comment

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

The class name is different in praxis and flax.

Any reason for that?
@mingxu1067mingxu1067May 2, 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.

FusedSoftmax is more sutiable to its behavior. We had a thought to change the name in Flax. But haven't have a conclusion yet.
@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
c4fb418 to
dc4dfc1CompareMay 2, 2023 03:22
@mingxu1067
mingxu1067
commented
May 2, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Update:
  1. Adapt to the latest Praxis.
@timmoon10
timmoon10
self-requested a review
May 3, 2023 23:12
timmoon10
timmoon10
approved these changes
May 4, 2023
View reviewed changes
@timmoon10timmoon10
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 once we fix the test failures.
@nouiz
nouiz
commented
May 8, 2023•

edited

Loading

Copy link


Copy Markdown

Collaborator
@mingxu1067 as we only need to fix the test failure, can you remove the draft status?
@mingxu1067
mingxu1067
commented
May 9, 2023•

edited

Loading

Copy link


Copy Markdown

CollaboratorAuthor
@nouiz , I will rebase this on top of #196 then mark as ready.
@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
dc4dfc1 to
116ce7aCompareMay 10, 2023 02:03
@mingxu1067
mingxu1067
marked this pull request as ready for review
May 10, 2023 02:05
@mingxu1067
mingxu1067
commented
May 10, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Rebased on main.
timmoon10
timmoon10
requested changes
May 12, 2023
View reviewed changes
Comment threadtransformer_engine/jax/__init__.py
Outdated

@mingxu1067
Adding JAX/Praxis modules and dependencies.
b75b0b9
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Adding UTs to JAX/Praxis modules.
23410e2
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Remove praxis as a dependency due to not strictly needed
08d09ab
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Repalce is_fp8_supported to is_fp8_available
11a8e66
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Make Praxis as an optional dependency.
4823c01
1. Removed 'from . import praxis' in __init__.py.
1.1 Noted, keep 'from . import flax' for deprecated warning.
2. Changed te.flax to te_flax in examples and README.rst.
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
cadc4b0 to
4823c01CompareMay 15, 2023 02:53
@mingxu1067
mingxu1067
commented
May 15, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@mingxu1067mingxu1067
changed the title
The Implementation of Praxis's Modules for Pre-reviewThe Implementation of Praxis's ModulesMay 15, 2023
@mingxu1067
Adding a workaround to FP8 training on Praxis.
c72007e
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
mingxu1067
commented
May 15, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@mingxu1067
mingxu1067
commented
May 15, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
timmoon10
timmoon10
approved these changes
May 15, 2023
View reviewed changes
@timmoon10timmoon10
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
@mingxu1067
mingxu1067
commented
May 16, 2023

Copy link


Copy Markdown

CollaboratorAuthor
CI is passed, could we merge this? Thanks.
@timmoon10
timmoon10
merged commit b20c053
into
NVIDIA:mainMay 16, 2023
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@timmoon10timmoon10timmoon10 approved these changes

+1 more reviewer

@nouiznouiznouiz left review comments
Reviewers whose approvals may not affect merge requirements

Assignees


No one assigned


Labels


None yet


Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

3 participants

@mingxu1067@nouiz@timmoon10

, '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('^' + ".*" + '
The Implementation of Praxis's Modules by mingxu1067 · Pull Request #158 · NVIDIA/TransformerEngine · GitHub
Skip to content
The Implementation of Praxis's Modules - #158
Merged
timmoon10 merged 7 commits into
NVIDIA:mainfrom 
mingxu1067:mingh/praxis_modules_for_review
May 16, 2023
Merged
The Implementation of Praxis's Modules#158
timmoon10 merged 7 commits into
NVIDIA:mainfrom 
mingxu1067:mingh/praxis_modules_for_review
Conversation
@mingxu1067
@mingxu1067mingxu1067
commented
Apr 19, 2023

Copy link


Copy Markdown

Collaborator
 No description provided. 
@mingxu1067
mingxu1067
commented
Apr 19, 2023

Copy link


Copy Markdown

CollaboratorAuthor
@nouiz and @timmoon10 , this is our initial implementation to support Praxis, and this is built on top of #152 .
@mingxu1067
mingxu1067
commented
Apr 19, 2023

Copy link


Copy Markdown

CollaboratorAuthor
@jeng1220 for viz.
@ksivaman
ksivaman
marked this pull request as draft
April 19, 2023 03:56
@ksivamanksivaman
changed the title
[WIP, DO NOT MERGE] The Implementation of Praxis's Modules for Pre-reviewThe Implementation of Praxis's Modules for Pre-reviewApr 19, 2023
nouiz
nouiz
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/jax/praxis/module.py
Outdated

@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
db8114d to
910948bCompareApril 21, 2023 03:22
@mingxu1067
mingxu1067
commented
Apr 21, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Update:
  1. Rebased on main.
  2. Added Zero-Centered-Gamma supports to LayerNorm.
nouiz
nouiz
reviewed
Apr 21, 2023
View reviewed changes
Comment threadtransformer_engine/jax/praxis/transformer.py
Outdated

@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
910948b to
c4fb418CompareApril 24, 2023 05:50
@mingxu1067
mingxu1067
commented
Apr 24, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Update:
  1. Rebased on [JAX] Adjust Module Structure. #169 .
  2. Added Zero-Centered-Gamma supports to all LN-related modules.
timmoon10
timmoon10
reviewed
Apr 24, 2023
View reviewed changes
@timmoon10timmoon10
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.

This PR looks reasonable to me, pending the merge of #169 and fixing test failures.
Comment threadtests/jax/test_praxis_layers.py
Outdated

@timmoon10timmoon10Apr 24, 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.

These tests are fine, but they could be made more challenging by comparing the Praxis layers against the non-TE implementations in https://github.com/NVIDIA/TransformerEngine/blob/main/tests/jax/utils.py. After all, the Praxis layers are mostly generated from the Flax layers anyways. That said, this is not high priority.
@mingxu1067mingxu1067May 2, 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.

Agree, put into TODO.
nouiz
nouiz
reviewed
May 1, 2023
View reviewed changes
Comment threadtransformer_engine/jax/praxis/module.py
Outdated

@nouiznouizMay 1, 2023

Copy link


Copy Markdown

Collaborator
Choose a reason for hiding this comment

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

The class name is different in praxis and flax.

Any reason for that?
@mingxu1067mingxu1067May 2, 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.

FusedSoftmax is more sutiable to its behavior. We had a thought to change the name in Flax. But haven't have a conclusion yet.
@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
c4fb418 to
dc4dfc1CompareMay 2, 2023 03:22
@mingxu1067
mingxu1067
commented
May 2, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Update:
  1. Adapt to the latest Praxis.
@timmoon10
timmoon10
self-requested a review
May 3, 2023 23:12
timmoon10
timmoon10
approved these changes
May 4, 2023
View reviewed changes
@timmoon10timmoon10
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 once we fix the test failures.
@nouiz
nouiz
commented
May 8, 2023•

edited

Loading

Copy link


Copy Markdown

Collaborator
@mingxu1067 as we only need to fix the test failure, can you remove the draft status?
@mingxu1067
mingxu1067
commented
May 9, 2023•

edited

Loading

Copy link


Copy Markdown

CollaboratorAuthor
@nouiz , I will rebase this on top of #196 then mark as ready.
@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
dc4dfc1 to
116ce7aCompareMay 10, 2023 02:03
@mingxu1067
mingxu1067
marked this pull request as ready for review
May 10, 2023 02:05
@mingxu1067
mingxu1067
commented
May 10, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Rebased on main.
timmoon10
timmoon10
requested changes
May 12, 2023
View reviewed changes
Comment threadtransformer_engine/jax/__init__.py
Outdated

@mingxu1067
Adding JAX/Praxis modules and dependencies.
b75b0b9
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Adding UTs to JAX/Praxis modules.
23410e2
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Remove praxis as a dependency due to not strictly needed
08d09ab
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Repalce is_fp8_supported to is_fp8_available
11a8e66
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Make Praxis as an optional dependency.
4823c01
1. Removed 'from . import praxis' in __init__.py.
1.1 Noted, keep 'from . import flax' for deprecated warning.
2. Changed te.flax to te_flax in examples and README.rst.
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
cadc4b0 to
4823c01CompareMay 15, 2023 02:53
@mingxu1067
mingxu1067
commented
May 15, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@mingxu1067mingxu1067
changed the title
The Implementation of Praxis's Modules for Pre-reviewThe Implementation of Praxis's ModulesMay 15, 2023
@mingxu1067
Adding a workaround to FP8 training on Praxis.
c72007e
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
mingxu1067
commented
May 15, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@mingxu1067
mingxu1067
commented
May 15, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
timmoon10
timmoon10
approved these changes
May 15, 2023
View reviewed changes
@timmoon10timmoon10
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
@mingxu1067
mingxu1067
commented
May 16, 2023

Copy link


Copy Markdown

CollaboratorAuthor
CI is passed, could we merge this? Thanks.
@timmoon10
timmoon10
merged commit b20c053
into
NVIDIA:mainMay 16, 2023
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@timmoon10timmoon10timmoon10 approved these changes

+1 more reviewer

@nouiznouiznouiz left review comments
Reviewers whose approvals may not affect merge requirements

Assignees


No one assigned


Labels


None yet


Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

3 participants

@mingxu1067@nouiz@timmoon10

, '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" + '
The Implementation of Praxis's Modules by mingxu1067 · Pull Request #158 · NVIDIA/TransformerEngine · GitHub
Skip to content
The Implementation of Praxis's Modules - #158
Merged
timmoon10 merged 7 commits into
NVIDIA:mainfrom 
mingxu1067:mingh/praxis_modules_for_review
May 16, 2023
Merged
The Implementation of Praxis's Modules#158
timmoon10 merged 7 commits into
NVIDIA:mainfrom 
mingxu1067:mingh/praxis_modules_for_review
Conversation
@mingxu1067
@mingxu1067mingxu1067
commented
Apr 19, 2023

Copy link


Copy Markdown

Collaborator
 No description provided. 
@mingxu1067
mingxu1067
commented
Apr 19, 2023

Copy link


Copy Markdown

CollaboratorAuthor
@nouiz and @timmoon10 , this is our initial implementation to support Praxis, and this is built on top of #152 .
@mingxu1067
mingxu1067
commented
Apr 19, 2023

Copy link


Copy Markdown

CollaboratorAuthor
@jeng1220 for viz.
@ksivaman
ksivaman
marked this pull request as draft
April 19, 2023 03:56
@ksivamanksivaman
changed the title
[WIP, DO NOT MERGE] The Implementation of Praxis's Modules for Pre-reviewThe Implementation of Praxis's Modules for Pre-reviewApr 19, 2023
nouiz
nouiz
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/jax/praxis/module.py
Outdated

@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
db8114d to
910948bCompareApril 21, 2023 03:22
@mingxu1067
mingxu1067
commented
Apr 21, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Update:
  1. Rebased on main.
  2. Added Zero-Centered-Gamma supports to LayerNorm.
nouiz
nouiz
reviewed
Apr 21, 2023
View reviewed changes
Comment threadtransformer_engine/jax/praxis/transformer.py
Outdated

@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
910948b to
c4fb418CompareApril 24, 2023 05:50
@mingxu1067
mingxu1067
commented
Apr 24, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Update:
  1. Rebased on [JAX] Adjust Module Structure. #169 .
  2. Added Zero-Centered-Gamma supports to all LN-related modules.
timmoon10
timmoon10
reviewed
Apr 24, 2023
View reviewed changes
@timmoon10timmoon10
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.

This PR looks reasonable to me, pending the merge of #169 and fixing test failures.
Comment threadtests/jax/test_praxis_layers.py
Outdated

@timmoon10timmoon10Apr 24, 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.

These tests are fine, but they could be made more challenging by comparing the Praxis layers against the non-TE implementations in https://github.com/NVIDIA/TransformerEngine/blob/main/tests/jax/utils.py. After all, the Praxis layers are mostly generated from the Flax layers anyways. That said, this is not high priority.
@mingxu1067mingxu1067May 2, 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.

Agree, put into TODO.
nouiz
nouiz
reviewed
May 1, 2023
View reviewed changes
Comment threadtransformer_engine/jax/praxis/module.py
Outdated

@nouiznouizMay 1, 2023

Copy link


Copy Markdown

Collaborator
Choose a reason for hiding this comment

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

The class name is different in praxis and flax.

Any reason for that?
@mingxu1067mingxu1067May 2, 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.

FusedSoftmax is more sutiable to its behavior. We had a thought to change the name in Flax. But haven't have a conclusion yet.
@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
c4fb418 to
dc4dfc1CompareMay 2, 2023 03:22
@mingxu1067
mingxu1067
commented
May 2, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Update:
  1. Adapt to the latest Praxis.
@timmoon10
timmoon10
self-requested a review
May 3, 2023 23:12
timmoon10
timmoon10
approved these changes
May 4, 2023
View reviewed changes
@timmoon10timmoon10
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 once we fix the test failures.
@nouiz
nouiz
commented
May 8, 2023•

edited

Loading

Copy link


Copy Markdown

Collaborator
@mingxu1067 as we only need to fix the test failure, can you remove the draft status?
@mingxu1067
mingxu1067
commented
May 9, 2023•

edited

Loading

Copy link


Copy Markdown

CollaboratorAuthor
@nouiz , I will rebase this on top of #196 then mark as ready.
@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
dc4dfc1 to
116ce7aCompareMay 10, 2023 02:03
@mingxu1067
mingxu1067
marked this pull request as ready for review
May 10, 2023 02:05
@mingxu1067
mingxu1067
commented
May 10, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Rebased on main.
timmoon10
timmoon10
requested changes
May 12, 2023
View reviewed changes
Comment threadtransformer_engine/jax/__init__.py
Outdated

@mingxu1067
Adding JAX/Praxis modules and dependencies.
b75b0b9
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Adding UTs to JAX/Praxis modules.
23410e2
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Remove praxis as a dependency due to not strictly needed
08d09ab
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Repalce is_fp8_supported to is_fp8_available
11a8e66
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Make Praxis as an optional dependency.
4823c01
1. Removed 'from . import praxis' in __init__.py.
1.1 Noted, keep 'from . import flax' for deprecated warning.
2. Changed te.flax to te_flax in examples and README.rst.
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
cadc4b0 to
4823c01CompareMay 15, 2023 02:53
@mingxu1067
mingxu1067
commented
May 15, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@mingxu1067mingxu1067
changed the title
The Implementation of Praxis's Modules for Pre-reviewThe Implementation of Praxis's ModulesMay 15, 2023
@mingxu1067
Adding a workaround to FP8 training on Praxis.
c72007e
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
mingxu1067
commented
May 15, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@mingxu1067
mingxu1067
commented
May 15, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
timmoon10
timmoon10
approved these changes
May 15, 2023
View reviewed changes
@timmoon10timmoon10
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
@mingxu1067
mingxu1067
commented
May 16, 2023

Copy link


Copy Markdown

CollaboratorAuthor
CI is passed, could we merge this? Thanks.
@timmoon10
timmoon10
merged commit b20c053
into
NVIDIA:mainMay 16, 2023
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@timmoon10timmoon10timmoon10 approved these changes

+1 more reviewer

@nouiznouiznouiz left review comments
Reviewers whose approvals may not affect merge requirements

Assignees


No one assigned


Labels


None yet


Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

3 participants

@mingxu1067@nouiz@timmoon10

, '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('^' + ".*" + '
The Implementation of Praxis's Modules by mingxu1067 · Pull Request #158 · NVIDIA/TransformerEngine · GitHub
Skip to content
The Implementation of Praxis's Modules - #158
Merged
timmoon10 merged 7 commits into
NVIDIA:mainfrom 
mingxu1067:mingh/praxis_modules_for_review
May 16, 2023
Merged
The Implementation of Praxis's Modules#158
timmoon10 merged 7 commits into
NVIDIA:mainfrom 
mingxu1067:mingh/praxis_modules_for_review
Conversation
@mingxu1067
@mingxu1067mingxu1067
commented
Apr 19, 2023

Copy link


Copy Markdown

Collaborator
 No description provided. 
@mingxu1067
mingxu1067
commented
Apr 19, 2023

Copy link


Copy Markdown

CollaboratorAuthor
@nouiz and @timmoon10 , this is our initial implementation to support Praxis, and this is built on top of #152 .
@mingxu1067
mingxu1067
commented
Apr 19, 2023

Copy link


Copy Markdown

CollaboratorAuthor
@jeng1220 for viz.
@ksivaman
ksivaman
marked this pull request as draft
April 19, 2023 03:56
@ksivamanksivaman
changed the title
[WIP, DO NOT MERGE] The Implementation of Praxis's Modules for Pre-reviewThe Implementation of Praxis's Modules for Pre-reviewApr 19, 2023
nouiz
nouiz
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/jax/praxis/module.py
Outdated

@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
db8114d to
910948bCompareApril 21, 2023 03:22
@mingxu1067
mingxu1067
commented
Apr 21, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Update:
  1. Rebased on main.
  2. Added Zero-Centered-Gamma supports to LayerNorm.
nouiz
nouiz
reviewed
Apr 21, 2023
View reviewed changes
Comment threadtransformer_engine/jax/praxis/transformer.py
Outdated

@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
910948b to
c4fb418CompareApril 24, 2023 05:50
@mingxu1067
mingxu1067
commented
Apr 24, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Update:
  1. Rebased on [JAX] Adjust Module Structure. #169 .
  2. Added Zero-Centered-Gamma supports to all LN-related modules.
timmoon10
timmoon10
reviewed
Apr 24, 2023
View reviewed changes
@timmoon10timmoon10
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.

This PR looks reasonable to me, pending the merge of #169 and fixing test failures.
Comment threadtests/jax/test_praxis_layers.py
Outdated

@timmoon10timmoon10Apr 24, 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.

These tests are fine, but they could be made more challenging by comparing the Praxis layers against the non-TE implementations in https://github.com/NVIDIA/TransformerEngine/blob/main/tests/jax/utils.py. After all, the Praxis layers are mostly generated from the Flax layers anyways. That said, this is not high priority.
@mingxu1067mingxu1067May 2, 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.

Agree, put into TODO.
nouiz
nouiz
reviewed
May 1, 2023
View reviewed changes
Comment threadtransformer_engine/jax/praxis/module.py
Outdated

@nouiznouizMay 1, 2023

Copy link


Copy Markdown

Collaborator
Choose a reason for hiding this comment

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

The class name is different in praxis and flax.

Any reason for that?
@mingxu1067mingxu1067May 2, 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.

FusedSoftmax is more sutiable to its behavior. We had a thought to change the name in Flax. But haven't have a conclusion yet.
@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
c4fb418 to
dc4dfc1CompareMay 2, 2023 03:22
@mingxu1067
mingxu1067
commented
May 2, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Update:
  1. Adapt to the latest Praxis.
@timmoon10
timmoon10
self-requested a review
May 3, 2023 23:12
timmoon10
timmoon10
approved these changes
May 4, 2023
View reviewed changes
@timmoon10timmoon10
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 once we fix the test failures.
@nouiz
nouiz
commented
May 8, 2023•

edited

Loading

Copy link


Copy Markdown

Collaborator
@mingxu1067 as we only need to fix the test failure, can you remove the draft status?
@mingxu1067
mingxu1067
commented
May 9, 2023•

edited

Loading

Copy link


Copy Markdown

CollaboratorAuthor
@nouiz , I will rebase this on top of #196 then mark as ready.
@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
dc4dfc1 to
116ce7aCompareMay 10, 2023 02:03
@mingxu1067
mingxu1067
marked this pull request as ready for review
May 10, 2023 02:05
@mingxu1067
mingxu1067
commented
May 10, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Rebased on main.
timmoon10
timmoon10
requested changes
May 12, 2023
View reviewed changes
Comment threadtransformer_engine/jax/__init__.py
Outdated

@mingxu1067
Adding JAX/Praxis modules and dependencies.
b75b0b9
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Adding UTs to JAX/Praxis modules.
23410e2
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Remove praxis as a dependency due to not strictly needed
08d09ab
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Repalce is_fp8_supported to is_fp8_available
11a8e66
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Make Praxis as an optional dependency.
4823c01
1. Removed 'from . import praxis' in __init__.py.
1.1 Noted, keep 'from . import flax' for deprecated warning.
2. Changed te.flax to te_flax in examples and README.rst.
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
cadc4b0 to
4823c01CompareMay 15, 2023 02:53
@mingxu1067
mingxu1067
commented
May 15, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@mingxu1067mingxu1067
changed the title
The Implementation of Praxis's Modules for Pre-reviewThe Implementation of Praxis's ModulesMay 15, 2023
@mingxu1067
Adding a workaround to FP8 training on Praxis.
c72007e
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
mingxu1067
commented
May 15, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@mingxu1067
mingxu1067
commented
May 15, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
timmoon10
timmoon10
approved these changes
May 15, 2023
View reviewed changes
@timmoon10timmoon10
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
@mingxu1067
mingxu1067
commented
May 16, 2023

Copy link


Copy Markdown

CollaboratorAuthor
CI is passed, could we merge this? Thanks.
@timmoon10
timmoon10
merged commit b20c053
into
NVIDIA:mainMay 16, 2023
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@timmoon10timmoon10timmoon10 approved these changes

+1 more reviewer

@nouiznouiznouiz left review comments
Reviewers whose approvals may not affect merge requirements

Assignees


No one assigned


Labels


None yet


Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

3 participants

@mingxu1067@nouiz@timmoon10

, '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('^' + ".*" + '
The Implementation of Praxis's Modules by mingxu1067 · Pull Request #158 · NVIDIA/TransformerEngine · GitHub
Skip to content
The Implementation of Praxis's Modules - #158
Merged
timmoon10 merged 7 commits into
NVIDIA:mainfrom 
mingxu1067:mingh/praxis_modules_for_review
May 16, 2023
Merged
The Implementation of Praxis's Modules#158
timmoon10 merged 7 commits into
NVIDIA:mainfrom 
mingxu1067:mingh/praxis_modules_for_review
Conversation
@mingxu1067
@mingxu1067mingxu1067
commented
Apr 19, 2023

Copy link


Copy Markdown

Collaborator
 No description provided. 
@mingxu1067
mingxu1067
commented
Apr 19, 2023

Copy link


Copy Markdown

CollaboratorAuthor
@nouiz and @timmoon10 , this is our initial implementation to support Praxis, and this is built on top of #152 .
@mingxu1067
mingxu1067
commented
Apr 19, 2023

Copy link


Copy Markdown

CollaboratorAuthor
@jeng1220 for viz.
@ksivaman
ksivaman
marked this pull request as draft
April 19, 2023 03:56
@ksivamanksivaman
changed the title
[WIP, DO NOT MERGE] The Implementation of Praxis's Modules for Pre-reviewThe Implementation of Praxis's Modules for Pre-reviewApr 19, 2023
nouiz
nouiz
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/jax/praxis/module.py
Outdated

@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
db8114d to
910948bCompareApril 21, 2023 03:22
@mingxu1067
mingxu1067
commented
Apr 21, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Update:
  1. Rebased on main.
  2. Added Zero-Centered-Gamma supports to LayerNorm.
nouiz
nouiz
reviewed
Apr 21, 2023
View reviewed changes
Comment threadtransformer_engine/jax/praxis/transformer.py
Outdated

@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
910948b to
c4fb418CompareApril 24, 2023 05:50
@mingxu1067
mingxu1067
commented
Apr 24, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Update:
  1. Rebased on [JAX] Adjust Module Structure. #169 .
  2. Added Zero-Centered-Gamma supports to all LN-related modules.
timmoon10
timmoon10
reviewed
Apr 24, 2023
View reviewed changes
@timmoon10timmoon10
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.

This PR looks reasonable to me, pending the merge of #169 and fixing test failures.
Comment threadtests/jax/test_praxis_layers.py
Outdated

@timmoon10timmoon10Apr 24, 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.

These tests are fine, but they could be made more challenging by comparing the Praxis layers against the non-TE implementations in https://github.com/NVIDIA/TransformerEngine/blob/main/tests/jax/utils.py. After all, the Praxis layers are mostly generated from the Flax layers anyways. That said, this is not high priority.
@mingxu1067mingxu1067May 2, 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.

Agree, put into TODO.
nouiz
nouiz
reviewed
May 1, 2023
View reviewed changes
Comment threadtransformer_engine/jax/praxis/module.py
Outdated

@nouiznouizMay 1, 2023

Copy link


Copy Markdown

Collaborator
Choose a reason for hiding this comment

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

The class name is different in praxis and flax.

Any reason for that?
@mingxu1067mingxu1067May 2, 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.

FusedSoftmax is more sutiable to its behavior. We had a thought to change the name in Flax. But haven't have a conclusion yet.
@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
c4fb418 to
dc4dfc1CompareMay 2, 2023 03:22
@mingxu1067
mingxu1067
commented
May 2, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Update:
  1. Adapt to the latest Praxis.
@timmoon10
timmoon10
self-requested a review
May 3, 2023 23:12
timmoon10
timmoon10
approved these changes
May 4, 2023
View reviewed changes
@timmoon10timmoon10
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 once we fix the test failures.
@nouiz
nouiz
commented
May 8, 2023•

edited

Loading

Copy link


Copy Markdown

Collaborator
@mingxu1067 as we only need to fix the test failure, can you remove the draft status?
@mingxu1067
mingxu1067
commented
May 9, 2023•

edited

Loading

Copy link


Copy Markdown

CollaboratorAuthor
@nouiz , I will rebase this on top of #196 then mark as ready.
@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
dc4dfc1 to
116ce7aCompareMay 10, 2023 02:03
@mingxu1067
mingxu1067
marked this pull request as ready for review
May 10, 2023 02:05
@mingxu1067
mingxu1067
commented
May 10, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Rebased on main.
timmoon10
timmoon10
requested changes
May 12, 2023
View reviewed changes
Comment threadtransformer_engine/jax/__init__.py
Outdated

@mingxu1067
Adding JAX/Praxis modules and dependencies.
b75b0b9
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Adding UTs to JAX/Praxis modules.
23410e2
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Remove praxis as a dependency due to not strictly needed
08d09ab
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Repalce is_fp8_supported to is_fp8_available
11a8e66
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Make Praxis as an optional dependency.
4823c01
1. Removed 'from . import praxis' in __init__.py.
1.1 Noted, keep 'from . import flax' for deprecated warning.
2. Changed te.flax to te_flax in examples and README.rst.
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
cadc4b0 to
4823c01CompareMay 15, 2023 02:53
@mingxu1067
mingxu1067
commented
May 15, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@mingxu1067mingxu1067
changed the title
The Implementation of Praxis's Modules for Pre-reviewThe Implementation of Praxis's ModulesMay 15, 2023
@mingxu1067
Adding a workaround to FP8 training on Praxis.
c72007e
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
mingxu1067
commented
May 15, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@mingxu1067
mingxu1067
commented
May 15, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
timmoon10
timmoon10
approved these changes
May 15, 2023
View reviewed changes
@timmoon10timmoon10
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
@mingxu1067
mingxu1067
commented
May 16, 2023

Copy link


Copy Markdown

CollaboratorAuthor
CI is passed, could we merge this? Thanks.
@timmoon10
timmoon10
merged commit b20c053
into
NVIDIA:mainMay 16, 2023
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@timmoon10timmoon10timmoon10 approved these changes

+1 more reviewer

@nouiznouiznouiz left review comments
Reviewers whose approvals may not affect merge requirements

Assignees


No one assigned


Labels


None yet


Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

3 participants

@mingxu1067@nouiz@timmoon10

, '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); }
})();
})();
The Implementation of Praxis's Modules by mingxu1067 · Pull Request #158 · NVIDIA/TransformerEngine · GitHub
Skip to content
The Implementation of Praxis's Modules - #158
Merged
timmoon10 merged 7 commits into
NVIDIA:mainfrom 
mingxu1067:mingh/praxis_modules_for_review
May 16, 2023
Merged
The Implementation of Praxis's Modules#158
timmoon10 merged 7 commits into
NVIDIA:mainfrom 
mingxu1067:mingh/praxis_modules_for_review
Conversation
@mingxu1067
@mingxu1067mingxu1067
commented
Apr 19, 2023

Copy link


Copy Markdown

Collaborator
 No description provided. 
@mingxu1067
mingxu1067
commented
Apr 19, 2023

Copy link


Copy Markdown

CollaboratorAuthor
@nouiz and @timmoon10 , this is our initial implementation to support Praxis, and this is built on top of #152 .
@mingxu1067
mingxu1067
commented
Apr 19, 2023

Copy link


Copy Markdown

CollaboratorAuthor
@jeng1220 for viz.
@ksivaman
ksivaman
marked this pull request as draft
April 19, 2023 03:56
@ksivamanksivaman
changed the title
[WIP, DO NOT MERGE] The Implementation of Praxis's Modules for Pre-reviewThe Implementation of Praxis's Modules for Pre-reviewApr 19, 2023
nouiz
nouiz
reviewed
Apr 20, 2023
View reviewed changes
Comment threadtransformer_engine/jax/praxis/module.py
Outdated

@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
db8114d to
910948bCompareApril 21, 2023 03:22
@mingxu1067
mingxu1067
commented
Apr 21, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Update:
  1. Rebased on main.
  2. Added Zero-Centered-Gamma supports to LayerNorm.
nouiz
nouiz
reviewed
Apr 21, 2023
View reviewed changes
Comment threadtransformer_engine/jax/praxis/transformer.py
Outdated

@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
910948b to
c4fb418CompareApril 24, 2023 05:50
@mingxu1067
mingxu1067
commented
Apr 24, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Update:
  1. Rebased on [JAX] Adjust Module Structure. #169 .
  2. Added Zero-Centered-Gamma supports to all LN-related modules.
timmoon10
timmoon10
reviewed
Apr 24, 2023
View reviewed changes
@timmoon10timmoon10
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.

This PR looks reasonable to me, pending the merge of #169 and fixing test failures.
Comment threadtests/jax/test_praxis_layers.py
Outdated

@timmoon10timmoon10Apr 24, 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.

These tests are fine, but they could be made more challenging by comparing the Praxis layers against the non-TE implementations in https://github.com/NVIDIA/TransformerEngine/blob/main/tests/jax/utils.py. After all, the Praxis layers are mostly generated from the Flax layers anyways. That said, this is not high priority.
@mingxu1067mingxu1067May 2, 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.

Agree, put into TODO.
nouiz
nouiz
reviewed
May 1, 2023
View reviewed changes
Comment threadtransformer_engine/jax/praxis/module.py
Outdated

@nouiznouizMay 1, 2023

Copy link


Copy Markdown

Collaborator
Choose a reason for hiding this comment

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

The class name is different in praxis and flax.

Any reason for that?
@mingxu1067mingxu1067May 2, 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.

FusedSoftmax is more sutiable to its behavior. We had a thought to change the name in Flax. But haven't have a conclusion yet.
@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
c4fb418 to
dc4dfc1CompareMay 2, 2023 03:22
@mingxu1067
mingxu1067
commented
May 2, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Update:
  1. Adapt to the latest Praxis.
@timmoon10
timmoon10
self-requested a review
May 3, 2023 23:12
timmoon10
timmoon10
approved these changes
May 4, 2023
View reviewed changes
@timmoon10timmoon10
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 once we fix the test failures.
@nouiz
nouiz
commented
May 8, 2023•

edited

Loading

Copy link


Copy Markdown

Collaborator
@mingxu1067 as we only need to fix the test failure, can you remove the draft status?
@mingxu1067
mingxu1067
commented
May 9, 2023•

edited

Loading

Copy link


Copy Markdown

CollaboratorAuthor
@nouiz , I will rebase this on top of #196 then mark as ready.
@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
dc4dfc1 to
116ce7aCompareMay 10, 2023 02:03
@mingxu1067
mingxu1067
marked this pull request as ready for review
May 10, 2023 02:05
@mingxu1067
mingxu1067
commented
May 10, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Rebased on main.
timmoon10
timmoon10
requested changes
May 12, 2023
View reviewed changes
Comment threadtransformer_engine/jax/__init__.py
Outdated

@mingxu1067
Adding JAX/Praxis modules and dependencies.
b75b0b9
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Adding UTs to JAX/Praxis modules.
23410e2
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Remove praxis as a dependency due to not strictly needed
08d09ab
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Repalce is_fp8_supported to is_fp8_available
11a8e66
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Make Praxis as an optional dependency.
4823c01
1. Removed 'from . import praxis' in __init__.py.
1.1 Noted, keep 'from . import flax' for deprecated warning.
2. Changed te.flax to te_flax in examples and README.rst.
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
mingxu1067force-pushed
the
mingh/praxis_modules_for_review
branch
from
cadc4b0 to
4823c01CompareMay 15, 2023 02:53
@mingxu1067
mingxu1067
commented
May 15, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@mingxu1067mingxu1067
changed the title
The Implementation of Praxis's Modules for Pre-reviewThe Implementation of Praxis's ModulesMay 15, 2023
@mingxu1067
Adding a workaround to FP8 training on Praxis.
c72007e
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
mingxu1067
commented
May 15, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
@mingxu1067
mingxu1067
commented
May 15, 2023

Copy link


Copy Markdown

CollaboratorAuthor
/te-ci
timmoon10
timmoon10
approved these changes
May 15, 2023
View reviewed changes
@timmoon10timmoon10
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
@mingxu1067
mingxu1067
commented
May 16, 2023

Copy link


Copy Markdown

CollaboratorAuthor
CI is passed, could we merge this? Thanks.
@timmoon10
timmoon10
merged commit b20c053
into
NVIDIA:mainMay 16, 2023
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@timmoon10timmoon10timmoon10 approved these changes

+1 more reviewer

@nouiznouiznouiz left review comments
Reviewers whose approvals may not affect merge requirements

Assignees


No one assigned


Labels


None yet


Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

3 participants

@mingxu1067@nouiz@timmoon10