Skip to content

Fix Bugs of TE/JAX - #119

Merged
timmoon10 merged 7 commits into
NVIDIA:mainfrom
mingxu1067:mingh/fix_bugs
Mar 29, 2023
Merged

Fix Bugs of TE/JAX#119
timmoon10 merged 7 commits into
NVIDIA:mainfrom
mingxu1067:mingh/fix_bugs

Conversation

@mingxu1067

@mingxu1067mingxu1067 commented Mar 27, 2023

Copy link
Copy Markdown
Collaborator
  1. Added support to transpose_batch_seqlen when decoded=True.
  2. Fixed wrong masking issues in MultiHeadAttention when decoded=True.
  3. Fixed missing dropout broadcast dims in LayerNormMLP.

@mingxu1067

mingxu1067 commented Mar 27, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

1 similar comment
@mingxu1067

mingxu1067 commented Mar 27, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

@mingxu1067

mingxu1067 commented Mar 27, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

@mingxu1067mingxu1067 added bug Something isn't working v0.7.0 labels Mar 27, 2023
@mingxu1067

mingxu1067 commented Mar 27, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

@jeng1220

jeng1220 commented Mar 27, 2023•
edited
Loading

Copy link
Copy Markdown
Contributor

@ptrendx and @timmoon10 ,
It is an important bugfix for T5X model training.

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

Seems reasonable to me, but I'd like to better understand the bug you're fixing.

Comment threadtransformer_engine/jax/transformer.py Outdated
@mingxu1067
Support transpose_bs when decoded=True
bf04878
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Fix Bugs,
3169db5
1. Fix missing dropout_dims in LayerNormMLP.
2. Fix broadcast issues in decoded.
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Fix wrong masks in decoded.
e5d6c5a
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Fixed wrong assert condition in TransformerLayer
6913860
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Fix amax is not set as 0 in each step.
bcffeed
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067

mingxu1067 commented Mar 28, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

@mingxu1067
Enhance rules conflict checking and docs.
15cd029
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067

mingxu1067 commented Mar 28, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

@mingxu1067
fix code formatting.
c4c6622
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067

mingxu1067 commented Mar 28, 2023

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

@timmoon10
timmoon10 merged commit f56e4fd into NVIDIA:mainMar 29, 2023
cyanguwa pushed a commit to cyanguwa/TransformerEngine that referenced this pull request Mar 31, 2023
@mingxu1067@cyanguwa
Fix Bugs of TE/JAX (NVIDIA#119)
32f44fe
* Support transpose_bs when decoded=True
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fix Bugs,
1. Fix missing dropout_dims in LayerNormMLP.
2. Fix broadcast issues in decoded.
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fix wrong masks in decoded.
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fixed wrong assert condition in TransformerLayer
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fix amax is not set as 0 in each step.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Enhance rules conflict checking and docs.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* fix code formatting.
Signed-off-by: Ming Huang <mingh@nvidia.com>
---------
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
ptrendx pushed a commit that referenced this pull request Mar 31, 2023
@mingxu1067@ptrendx
Fix Bugs of TE/JAX (#119)
bc0e448
* Support transpose_bs when decoded=True
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fix Bugs,
1. Fix missing dropout_dims in LayerNormMLP.
2. Fix broadcast issues in decoded.
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fix wrong masks in decoded.
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fixed wrong assert condition in TransformerLayer
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fix amax is not set as 0 in each step.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Enhance rules conflict checking and docs.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* fix code formatting.
Signed-off-by: Ming Huang <mingh@nvidia.com>
---------
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>
cyanguwa pushed a commit to cyanguwa/TransformerEngine that referenced this pull request Apr 1, 2023
@mingxu1067@cyanguwa
Fix Bugs of TE/JAX (NVIDIA#119)
c18bc89
* Support transpose_bs when decoded=True
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fix Bugs,
1. Fix missing dropout_dims in LayerNormMLP.
2. Fix broadcast issues in decoded.
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fix wrong masks in decoded.
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fixed wrong assert condition in TransformerLayer
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Fix amax is not set as 0 in each step.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* Enhance rules conflict checking and docs.
Signed-off-by: Ming Huang <mingh@nvidia.com>
* fix code formatting.
Signed-off-by: Ming Huang <mingh@nvidia.com>
---------
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
Signed-off-by: Ming Huang <mingh@nvidia.com>
Signed-off-by: Charlene Yang <charleney@nvidia.com>
@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

Assignees

No one assigned

Labels

bugSomething isn't workingv0.7.0

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@mingxu1067@jeng1220@timmoon10