') + ')', '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 TE examples by jeng1220 · Pull Request #108 · NVIDIA/TransformerEngine · GitHub
Skip to content

[JAX] Add TE examples - #108

Merged
timmoon10 merged 16 commits into
NVIDIA:mainfrom
jeng1220:rjeng/te_jax_examples
Mar 28, 2023
Merged

[JAX] Add TE examples#108
timmoon10 merged 16 commits into
NVIDIA:mainfrom
jeng1220:rjeng/te_jax_examples

Conversation

@jeng1220

@jeng1220jeng1220 commented Mar 20, 2023

Copy link
Copy Markdown
Contributor

Add following examples:

  • examples/jax/test_single_gpu_mnist.py
  • examples/jax/test_single_gpu_encoder.py
  • examples/jax/test_multigpu_encoder.py
  • examples/jax/test_model_parallel_encoder.py

Align default values to google/flax

  • transformer_engine/jax/module.py
  • transformer_engine/jax/transformer.py

@jeng1220jeng1220 changed the title [JAX] Add TE/JAX examples[JAX] Add TE examplesMar 20, 2023
@jeng1220jeng1220 mentioned this pull request Mar 20, 2023
Merged
nouiz
nouiz reviewed Mar 20, 2023
View reviewed changes
Comment threadtests/jax/test_mnist.py Outdated
Comment threadqa/L0_jax_unittest/test.sh Outdated
Comment threadexamples/jax/encoder/README.md Outdated
nouiz
nouiz reviewed Mar 20, 2023
View reviewed changes
Comment threadexamples/jax/encoder/test_multigpu_encoder.py Outdated
@timmoon10

timmoon10 commented Mar 21, 2023

Copy link
Copy Markdown
Member

/te-ci

@jeng1220
jeng1220force-pushed the rjeng/te_jax_examples branch 2 times, most recently from 6e9e380 to 53a11faCompareMarch 22, 2023 11:19
nouiz
nouiz reviewed Mar 22, 2023
View reviewed changes
Comment threadexamples/jax/mnist/README.md Outdated
Comment threadexamples/jax/mnist/README.md Outdated
Comment threadexamples/jax/mnist/requirements.txt Outdated
Comment threadtests/jax/test_mnist.py Outdated
Comment threadexamples/jax/mnist/test_single_gpu_mnist.py Outdated
Comment threadexamples/jax/encoder/README.md Outdated
Comment threadexamples/jax/encoder/test_multigpu_encoder.py Outdated
Comment threadexamples/jax/encoder/test_multigpu_encoder.py Outdated
Comment threadexamples/jax/encoder/test_multigpu_encoder.py Outdated
Comment threadexamples/jax/encoder/test_model_parallel_encoder.py Outdated
@jeng1220
refactor JAX examples
996b7cf
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
@jeng1220
fix doc-string
158c67b
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
@jeng1220
add dp example
acef7b0
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
@jeng1220
refactor
b2fa038
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
@jeng1220
fix params_axes_pspec
6d4fad7
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
@jeng1220
Add model parallel example and refactor
6d4a26a
Update readme
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
@jeng1220
align code and readme
dd9d4d3
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
@jeng1220
update verification
7b65d7f
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
@jeng1220
add mask
1727d7d
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
@jeng1220
num_gpu is configurable
d9cf72c
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
@jeng1220
update readme
b7869df
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
@jeng1220
jeng1220force-pushed the rjeng/te_jax_examples branch from 4ac12af to b7869dfCompareMarch 25, 2023 04:16
@jeng1220
update readme
773be64
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
@jeng1220
solvepylint issue
0455428
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
@jeng1220
ignore markdown and txt file from license check
aefbcad
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
@jeng1220
jeng1220force-pushed the rjeng/te_jax_examples branch from 94390f7 to aefbcadCompareMarch 25, 2023 04:58
@jeng1220

jeng1220 commented Mar 25, 2023•
edited
Loading

Copy link
Copy Markdown
ContributorAuthor

@timmoon10 and @nouiz ,
I think all conversation are resolved.
And I see internal pipeline 7708362 is good.

@nouiz

nouiz commented Mar 27, 2023

Copy link
Copy Markdown
Collaborator

@timmoon10 and @nouiz , I think all conversation are resolved. And I see internal pipeline 7708362 is good.

For me it is good except for one thing. GH doesn't allow me to make the comment at the right place, so trying here.

For the requirement, we still need to add flax as we don't always build and install in the same env.
It happens frequently that people use what is on pypi or build and one env and reuse in another one (We do this frequently in Dockerfile).

So we also need to add flax, even if in some cases it won't be useful.

@jeng1220
Update README.md
e14860c
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
@jeng1220
jeng1220force-pushed the rjeng/te_jax_examples branch from 3642ac0 to e14860cCompareMarch 27, 2023 13:18
@jeng1220
add flax into requirements.txt
8974af7
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
@jeng1220

jeng1220 commented Mar 27, 2023

Copy link
Copy Markdown
ContributorAuthor

For me it is good except for one thing. GH doesn't allow me to make the comment at the right place, so trying here.

For the requirement, we still need to add flax as we don't always build and install in the same env. It happens frequently that people use what is on pypi or build and one env and reuse in another one (We do this frequently in Dockerfile).

So we also need to add flax, even if in some cases it won't be useful.

Done

@mingxu1067

mingxu1067 commented Mar 28, 2023

Copy link
Copy Markdown
Collaborator

/te-ci

@nouiz

nouiz commented Mar 28, 2023

Copy link
Copy Markdown
Collaborator

What is the result of the CI?
What is blocking this PR from being merged?

timmoon10
timmoon10 approved these changes Mar 28, 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

The CI is waiting on some A100 tests, but everything else has passed.

@timmoon10
timmoon10 merged commit 5992e03 into NVIDIA:mainMar 28, 2023
@jeng1220

jeng1220 commented Mar 28, 2023

Copy link
Copy Markdown
ContributorAuthor

@timmoon10 and @nouiz ,
Thanks for help

cyanguwa pushed a commit to cyanguwa/TransformerEngine that referenced this pull request Mar 31, 2023
@jeng1220@cyanguwa
[JAX] Add TE examples (NVIDIA#108)
3c2eed4
* refactor JAX examples
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* fix doc-string
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* add dp example
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* refactor
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* fix params_axes_pspec
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* Add model parallel example and refactor
Update readme
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* align code and readme
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* update verification
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* add mask
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* num_gpu is configurable
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* update readme
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* update readme
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* solvepylint issue
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* ignore markdown and txt file from license check
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* Update README.md
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* add flax into requirements.txt
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
---------
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
ptrendx pushed a commit that referenced this pull request Mar 31, 2023
@jeng1220@ptrendx
[JAX] Add TE examples (#108)
084b1e5
* refactor JAX examples
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* fix doc-string
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* add dp example
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* refactor
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* fix params_axes_pspec
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* Add model parallel example and refactor
Update readme
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* align code and readme
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* update verification
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* add mask
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* num_gpu is configurable
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* update readme
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* update readme
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* solvepylint issue
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* ignore markdown and txt file from license check
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* Update README.md
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* add flax into requirements.txt
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
---------
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
cyanguwa pushed a commit to cyanguwa/TransformerEngine that referenced this pull request Apr 1, 2023
@jeng1220@cyanguwa
[JAX] Add TE examples (NVIDIA#108)
dd6001f
* refactor JAX examples
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* fix doc-string
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* add dp example
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* refactor
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* fix params_axes_pspec
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* Add model parallel example and refactor
Update readme
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* align code and readme
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* update verification
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* add mask
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* num_gpu is configurable
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* update readme
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* update readme
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* solvepylint issue
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* ignore markdown and txt file from license check
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* Update README.md
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
* add flax into requirements.txt
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
---------
Signed-off-by: Ryan Jeng <rjeng@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@jeng1220
jeng1220 deleted the rjeng/te_jax_examples branch April 24, 2023 01:45
@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

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

Assignees

No one assigned

Labels

v0.7.0

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@jeng1220@timmoon10@nouiz@mingxu1067@ptrendx